(self,
init_optimizer,
param_names,
bfloat16_config,
mpu=None,
clip_grad=0.0,
norm_type=2,
allgather_bucket_size=5000000000,
dp_process_group=None,
timers=None,
grad_acc_dtype=None,
graph_harvesting=False,
has_moe_layers=False)
| 37 | class BF16_Optimizer(ZeROOptimizer): |
| 38 | |
| 39 | def __init__(self, |
| 40 | init_optimizer, |
| 41 | param_names, |
| 42 | bfloat16_config, |
| 43 | mpu=None, |
| 44 | clip_grad=0.0, |
| 45 | norm_type=2, |
| 46 | allgather_bucket_size=5000000000, |
| 47 | dp_process_group=None, |
| 48 | timers=None, |
| 49 | grad_acc_dtype=None, |
| 50 | graph_harvesting=False, |
| 51 | has_moe_layers=False): |
| 52 | super().__init__() |
| 53 | see_memory_usage('begin bf16_optimizer', force=True) |
| 54 | self.timers = timers |
| 55 | self.optimizer = init_optimizer |
| 56 | self.param_names = param_names |
| 57 | self.using_real_optimizer = not isinstance(self.optimizer, DummyOptim) |
| 58 | |
| 59 | assert bfloat16_config.enabled, "BF16Optimizer: requires bfloat16 to be enabled" |
| 60 | assert grad_acc_dtype in [torch.float32, torch.bfloat16 |
| 61 | ], f"BF16Optimizer: Unsupported gradient accumulation data type: {grad_acc_dtype}" |
| 62 | self.grad_acc_dtype = grad_acc_dtype |
| 63 | |
| 64 | # BF16 doesn't use loss scaling, but these attributes are needed for API compatibility |
| 65 | self.custom_loss_scaler = False |
| 66 | self.external_loss_scale = None |
| 67 | self.torch_autocast_gradscaler = None |
| 68 | |
| 69 | self.immediate_grad_update = bfloat16_config.immediate_grad_update |
| 70 | |
| 71 | self.clip_grad = clip_grad |
| 72 | self.norm_type = norm_type |
| 73 | self.mpu = mpu |
| 74 | self.autoep_folding_tp_group = None |
| 75 | self.autoep_folding_spec = None |
| 76 | self.allgather_bucket_size = int(allgather_bucket_size) |
| 77 | self.dp_process_group = dp_process_group |
| 78 | self.dp_rank = dist.get_rank(group=self.dp_process_group) |
| 79 | self.has_moe_layers = has_moe_layers |
| 80 | self.non_expert_gradients = [] |
| 81 | self.real_dp_process_group = [dp_process_group for i in range(len(self.optimizer.param_groups))] |
| 82 | if self.has_moe_layers: |
| 83 | self._configure_moe_settings() |
| 84 | |
| 85 | # Use torch (un)flatten ops |
| 86 | self.flatten = _flatten_dense_tensors |
| 87 | self.unflatten = _unflatten_dense_tensors |
| 88 | |
| 89 | #align nccl all-gather send buffers to 4-bye boundary |
| 90 | self.nccl_start_alignment_factor = 2 # 4-byte alignment/sizeof(fp16) = 2 |
| 91 | |
| 92 | # Build BF16/FP32 groups |
| 93 | self.bf16_groups = [] |
| 94 | self.bf16_groups_flat = [] |
| 95 | self.bf16_partitioned_groups = [] |
| 96 |
nothing calls this directly
no test coverage detected