[FSDP][Loss] Reduce-sum gradient reduction with per-batch aux-loss hub#1968
[FSDP][Loss] Reduce-sum gradient reduction with per-batch aux-loss hub#1968HAOCHENYE wants to merge 1 commit into
Conversation
| step_grad_tokens = torch.zeros((), device=global_denominator.device) | ||
| for loss_ctx in loss_ctx_list: | ||
| step_grad_tokens += (loss_ctx.loss_kwargs.shifted_labels != loss_cfg.ignore_idx).sum() | ||
| display_coeff = (global_denominator / step_grad_tokens.clamp_min(1.0)).detach() |
There was a problem hiding this comment.
Shouldn't the coefficient here just be the global dominator divided by the rank dominator after the reduce sum? I think there is an issue with your current algorithm; calibrating at the token level this way is fine, but it doesn't work for sample-level and square-level losses.
| if z_loss is not None: | ||
| output["z_loss"] = z_loss | ||
| output["tokens_per_expert_global"] = tokens_per_expert_global | ||
| if aux_out["balancing_loss"] is not None: |
There was a problem hiding this comment.
I think we can update the contract for the model's output. The loss field in the output should represent the sum of all losses, with the engine layer handling simple gradient accumulation and the backward pass. Then, we can add a new loss_log field to store all the calibrated results, which the engine layer can simply sum up and include in the logs. Does that sound reasonable to you?
Also, there is currently a lot of hardcoding in post_micro_batch_forward, and the field names aren't very logical. My vision for the final format is: recording both local loss and reduced loss in TensorBoard, with logs appearing as loss(reduced): 1.2(1.3).
| # the reduce-scatter and are skipped (no Replicate placement). | ||
| self._reduce_replicated_ignored_grads(self.trainable_parameters()) | ||
|
|
||
| def _reduce_replicated_ignored_grads(self, params: Iterable[tuple[str, nn.Parameter]]) -> None: |
There was a problem hiding this comment.
I'm a bit confused here. Could you consider adding a comment? If I have TPSP enabled, an allgather is executed before the attention mechanism. If the attention structure includes QK normalization, that norm will process the full dataset. Wouldn't performing an allreduce sum on the QK norm gradients cause them to be scaled up by a factor of SP?
| summed: dict[str, torch.Tensor] = {} | ||
| for output in batch_outputs: | ||
| for name, value in (output.calibrated_losses or {}).items(): | ||
| contribution = value.detach().float() | ||
| summed[name] = contribution if name not in summed else summed[name] + contribution | ||
| return {f"reduced_{name}": value.item() for name, value in summed.items()} |
There was a problem hiding this comment.
Where does the reduction manifest? My understanding is that this is actually the local loss.
Switch FSDP gradient reduction from mean to pure SUM and remove the
compensating ×world_size injections across CE / balancing / z-loss, so each
parameter's gradient equals the global loss gradient with no compensating
divides.
- FSDP reduce-scatter -> SUM (set_gradient_divide_factor(1.0) +
set_force_sum_reduction_for_comms(True), avoids the bf16 PreMulSum zeroing);
scale_and_reduce_grad drops the expert div_(ep_size) and replicated
pre-divide and SUM-reduces Replicate-placement (incl. fp32-ignored) params.
- Loss: drop the CE WORLD all_reduce, balancing all_reduce_autograd, z-loss
×world_size, and the now-dead world_size plumbing.
- Display: per-rank calibrate() with no cross-rank all_reduce; each rank shows
its own per-token-mean loss.
- Aux losses folded into a per-batch AuxLossContext hub
(MoELossContextDict = {lm, aux, mtp}); AuxLossConfig.build(data, sp_mesh)
conforms to the loss-config contract; per-layer accumulate takes router-output
lists, applies padding removal / histc / balancing / z internally.
- Drop the non-global aux-loss averaging mode (incompatible with reduce-sum).
5f4f804 to
3ce8fea
Compare
HAOCHENYE
changed the title
1 participant
Switch FSDP gradient reduction from mean to pure SUM and remove the compensating ×world_size injections across CE / balancing / z-loss, so each parameter's gradient equals the global loss gradient with no compensating divides. Squashed to a single commit (the intermediate steps had incomplete calibration and would not individually pass CI).
set_gradient_divide_factor(1.0)+set_force_sum_reduction_for_comms(True)(avoids the bf16 NCCL PreMulSum zeroing), wired into everyfully_shardincl. compose/InternS1;scale_and_reduce_graddrops the expertdiv_(ep_size)/ replicated pre-divide and SUM-reducesReplicate-placement (incl. fp32-ignored) params.all_reduce, balancingall_reduce_autograd, z-loss× world_size, and the deadworld_sizeplumbing.calibrate()with no cross-rank all_reduce — each rank shows its own per-token-mean loss.AuxLossContexthub (MoELossContextDict = {lm, aux, mtp});AuxLossConfig.build(data, sp_mesh)conforms to the loss-config contract; per-layeraccumulatetakes router-output lists and does padding removal / histc / balancing / z internally.Tests:
tests/model/test_reduce_sum_grad.py(16) — bf16 SUM mechanism, EP1/EP2 token-mean parity, balancing & z-loss A/B, compose hook, Dense fp32 reduction, domino-MTP, SP world=4, per-rank display. pre-commit clean.Supersedes #1969 (folded in).