[sanketpurandare]

Prepared with AI assistance.

Reuse current main's optimizer-owned GradientClipper instead of adding engine
pre-update hooks or model norm iterators. Bind existing model_parts once through
the clipper's no-op extension hook after config-only construction, preserving
existing custom constructor compatibility. A model-specific implementation can finalize
shared gradients immediately before clipping, including optimizer CUDA graph
capture/replay. The default clipper's arithmetic stays unchanged.

Keep one model lifecycle seam for necessary post-initialization/checkpoint-load
weight synchronization. Extend distributed norm clipping with an optional
logical parameter subset while applying the same coefficient to every physical
copy. Both copies remain optimizer-owned; this does not drop replica updates.

Verification

Owning-commit CPU checks: 82 passed in 14.89s.

Strict relevant Sphinx builds and all non-type repository hooks passed.
Three unchanged main type findings remain attributed, not reported as a clean
full type run. External reviewer and CI acceptance are separate, still pending.

Complete-stack integration (not this PR in isolation): real four-GPU PP2/DP2/
EP2 and PP4/DP1/EP1 ran ten updates with exact model/optimizer tensor parity,
optimizer graphs and sequential validation. A fresh native PP2/PP4 checkpoint
matrix matched resumed state exactly. These are debug-model contract checks,
not 671B throughput or convergence results.

The final 207-case GPU matrix covers in-place gradients, dense attention, PP/FSDP train/eval graphs, shared-stream backward/capture and fused numerical contracts.

[tianyu-l]

seems too intrusive

[sanketpurandare]

Rebased onto main's optimizer-owned GradientClipper and removed the engine pre-update and model norm-iterator hooks. The remaining seams are init/load weight synchronization and a one-time binding hook that preserves existing config-only clipper constructors.

[tianyu-l]

@claude review

[claude]

Claude finished @tianyu-l's task in 2m 44s —— View job

---

Verdict: Request changes. The PR is less intrusive than before, but it can lose about half its core surface: norm_parameters can move into the DSv3 clipper, and the resync after checkpoint load seems redundant. I left 4 inline comments.

Details

What's justified

• GradientClipper.bind_model_parts is a small, generic, no-op seam. The embedding grads have to be finalized before the norm is computed. An optimizer step pre-hook runs too late, and _register_optimizer_hooks only receives OptimizersContainer, not the clipper. Keeping constructors config-only is reasonable.
• Synchronizing shared parameters after init is needed because set_determinism gives each PP rank a distinct seed, so the two embedding copies start out different.

What can be cut

1. norm_parameters in dist_utils.clip_grad_norm_ (torchtitan/distributed/utils.py:542). MTPGradientClipper ([DeepSeek V3] Add MTP1 eager pipeline parallelism #5329) can compute the norm over the logical parameters with the existing function. It can then call torch.nn.utils.clip_grads_with_norm_(replicas, max_norm, total_norm) with the same scaling factor. That gives the same result without touching a function every model uses. The new empty-norm branch also quietly changes the device of the zero norm when no parameter has a grad (now parameters[0].device instead of CPU).
2. The sync after checkpoint load (torchtitan/training_engine.py:928-934). Both copies keep the canonical tok_embeddings.weight FQN, so DCP already loads the same tensor into both. This hook also misses other load paths, such as torchft live recovery. Unless you can name a case where the copies differ after load, I'd drop it and keep only the sync after init.
3. Leftover test fixtures (tests/unit_tests/cpu/test_trainer.py:916-917, and possibly the model_cls default in _training_loop). optim_step doesn't read model_parts or parallelism_context in this PR.
4. README paragraph (torchtitan/config/README.md:87-95). It is mostly MTP-specific and repeats [DeepSeek V3] Add MTP1 eager pipeline parallelism #5329's MTP.md. One sentence about bind_model_parts is enough here.

General

• On its own, this PR adds three no-op seams; the only consumer is [DeepSeek V3] Add MTP1 eager pipeline parallelism #5329. Reviewers would find the change easier to judge if the PR description linked the consuming code: MTPGradientClipper, _synchronize_mtp_embeddings and _finalize_mtp_embedding_gradients.
• Other engines are fine. FaultTolerantTrainingEngine._initialize_model calls super(), and GraphTrainingEngine doesn't override init or load, so both still run the sync after init.
• I couldn't run the unit tests because the job's tool permissions don't allow pytest.

[claude]

Suggestion: I'd keep this function as it is and put the logical-norm logic in MTPGradientClipper (#5329). clip_grads_with_norm_ already lets you clip with a norm you computed yourself:

replicas = [p for p in parameters if id(p) in self._replica_ids]
logical = [p for p in parameters if id(p) not in self._replica_ids]
total_norm = dist_utils.clip_grad_norm_(
    logical, max_norm, foreach=True,
    pp_mesh=parallelism_context.get_optional_mesh("pp"),
)
torch.nn.utils.clip_grads_with_norm_(replicas, max_norm, total_norm, foreach=True)
return total_norm

That gives the same result: the replica is left out of the norm and both copies get the same scaling factor. It also removes about 27 lines from a function every model uses, along with the new empty-norm branch. That branch changes one existing behavior: when no parameter has a grad, the zero norm used to be created on CPU and is now created on parameters[0].device. Whether that change is good or bad, it shouldn't ship as a side effect of an MTP seam.

[claude]

Question: Why is the resync after load needed? According to #5329, both physical copies keep the canonical tok_embeddings.weight FQN. DCP should therefore load the same stored tensor into both owners, so the broadcast or copy here does nothing except add a collective on every resume.

This hook also doesn't cover every load path. For example, torchft live recovery restores state through the checkpointer without calling TrainingEngine.load_checkpoint. In practice, the canonical-FQN invariant is what keeps the copies in sync. If you know a case where the copies differ after load (partial load, exclude_from_loading, HF load through sd_adapter), please name it. Otherwise I'd drop this branch and keep only the sync after init. That sync is needed, because PP ranks get distinct seeds.

[claude]

Nit: TrainingEngine.optim_step doesn't read model_parts or parallelism_context in this PR. These look like leftovers from the engine pre-update hook you removed. The same applies to the model_cls = BaseModel default in _training_loop, unless one of those tests actually goes down the load_checkpoint() -> True path.

Suggested change

	model_parts=[],
	parallelism_context=SimpleNamespace(),

[claude]

Nit: This paragraph mostly describes MTP: shared replicas, finalizing grads before the norm, and CUDA graph replay. #5329's MTP.md already covers the same points almost word for word. In the general config README, I'd reduce it to one sentence, e.g. "Custom GradientClippers keep a config-only constructor; Optim calls bind_model_parts(model_parts) once after construction." Leave the rest to MTP.md.