(model, model_ema, beta=0.992, device=None, zero_stage=0)
| 153 | |
| 154 | |
| 155 | def moving_average(model, model_ema, beta=0.992, device=None, zero_stage=0): |
| 156 | zero_stage_3 = (zero_stage == 3) |
| 157 | with torch.no_grad(): |
| 158 | for param, param_ema in zip(model.parameters(), |
| 159 | model_ema.parameters()): |
| 160 | params_to_fetch = _z3_params_to_fetch([param, param_ema |
| 161 | ]) if zero_stage_3 else [] |
| 162 | should_gather_param = len(params_to_fetch) > 0 |
| 163 | with deepspeed.zero.GatheredParameters( |
| 164 | params_to_fetch, enabled=should_gather_param): |
| 165 | data = param.data |
| 166 | if device is not None: |
| 167 | data = data.to(device) |
| 168 | param_ema.data.copy_(torch.lerp(data, param_ema.data, beta)) |
| 169 | |
| 170 | |
| 171 | def save_zero_three_model(model_ema, global_rank, save_dir, zero_stage=0, sub_folder=""): |
nothing calls this directly
no test coverage detected