Mixed precision training

2-3배 빠른 training, 40% 적은 메모리, 거의 공짜

Mixed precision training은 compute에 낮은 정밀도 (float16/bfloat16) 쓰면서 weight 저장은 float32 유지. 최신 GPU랑 TPU에서 메모리 절약 + 처리량 향상, 정확도 영향 최소.

왜 작동하나: 신경망은 forward/backward pass의 정밀도 감소에 의외로 관대해. Float16 계산 gradient도 optimizer엔 충분히 정확. 하지만 weight는 float32 유지해야 작은 gradient 업데이트 (float16 epsilon보다 훨씬 작은)가 안 잃어버려.

Policy	Compute dtype	Variable dtype	적합한 곳
mixed_float16	float16	float32	NVIDIA GPU (V100, A100, H100)
mixed_bfloat16	bfloat16	float32	TPU (모든 버전)
float32	float32	float32	Default; 안정성 중요할 때

mixed_float16의 결정적 디테일: 마지막 activation은 수치 안정성 위해 float32. 출력 Activation layer (예: softmax)에 항상 dtype='float32' 설정. TPU (bfloat16)에선 bfloat16이 float32와 같은 exponent 범위 가져 불필요.

Code

Mixed precision policy 설정·python

import tensorflow as tf
from tensorflow.keras import mixed_precision, layers, Input, Model

# Set BEFORE building the model
mixed_precision.set_global_policy('mixed_float16')   # GPU
# mixed_precision.set_global_policy('mixed_bfloat16')  # TPU

policy = mixed_precision.Policy('mixed_float16')
print(f"Compute dtype: {policy.compute_dtype}")    # float16
print(f"Variable dtype: {policy.variable_dtype}")  # float32

# Build model — layers auto-use float16 for compute
inputs  = Input(shape=(784,))
x       = layers.Dense(4096, activation='relu')(inputs)
x       = layers.Dense(4096, activation='relu')(x)
x       = layers.Dense(10)(x)

# IMPORTANT: final activation must be float32 for numerical stability
outputs = layers.Activation('softmax', dtype='float32')(x)

model = Model(inputs, outputs)
model.compile(
    loss='sparse_categorical_crossentropy',
    optimizer=tf.keras.optimizers.Adam(),
    metrics=['accuracy'],
)

# Larger batches benefit more from float16 throughput
history = model.fit(x_train, y_train, batch_size=8192, epochs=5,
                    validation_split=0.2)

Custom loop + LossScaleOptimizer·python

import tensorflow as tf
from tensorflow.keras import mixed_precision

mixed_precision.set_global_policy('mixed_float16')

optimizer = tf.keras.optimizers.Adam()
optimizer = mixed_precision.LossScaleOptimizer(optimizer)  # prevent fp16 underflow

@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        predictions = model(x, training=True)
        loss = loss_fn(y, predictions)
        scaled_loss = optimizer.get_scaled_loss(loss)

    scaled_grads = tape.gradient(scaled_loss, model.trainable_variables)
    grads = optimizer.get_unscaled_gradients(scaled_grads)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

# Note: bfloat16 (TPU) does NOT need LossScaleOptimizer
# because bfloat16 has the same exponent range as float32

2-3배 빠른 training, 40% 적은 메모리, 거의 공짜

Code

External links

Progress

댓글 0