Mixed precision training in FP16

[xrsrke]

TODOs

• [ ] Come up with a design that does not modify the original modules like DataParallel, TensorParallel, PipelineParallel, ... in order to make them work in mixed precision training.
• [ ] Make 3D parallelism work in mixed precision training.
• [ ] Make DistributedOptimizer work in mixed precision training

APIs

import torch
import pipegoose

# other parallelism...
scaler = pipegoose.amp.GradScaler()

with pipegoose.amp.autocast(parallel_context, dtype=torch.float16):
    outputs = model(**inputs, labels=labels)
    loss = loss_func(outputs, targets)

scaled_loss = scaler.scale(loss)

optim.zero_grad()
scaled_loss.backward()
scaler.step(optimizer)
scaler.update() # updates the scale for next iteration

Reading List

• MixedPrecisionOptimizer [link] and Float16OptimizerWithFloat16Params [link] from Megatron-LM