fix: ep clipping with no ep grads

[garrett361]

The current EP grad clipping logic assumes that when using EP all of the norms returned by torch.nn.utils.get_total_norm are DTensors. This assumption can be violated and the subsequent full_tensor call can correspondingly fail in the edge case where the ep_grad list is empty, in which case get_total_norm returns tensor(0.), a non-DTensor.

torchtitan/torchtitan/distributed/utils.py

Lines 421 to 423 in a1fdd7e

	ep_grads_total_norm = torch.nn.utils.get_total_norm(
	ep_grads, norm_type, error_if_nonfinite, foreach
	).full_tensor()

File "/app/torchtitan/torchtitan/distributed/utils.py", line 423, in _clip_grad_norm_with_ep
    ).full_tensor()
      ^^^^^^^^^^^
AttributeError: 'Tensor' object has no attribute 'full_tensor'

This edge case can occur in PP+EP setups when model uses some fully dense and some MoE layers (like DSv3), in which case some pp ranks may not be assigned any MoE layers.

I suppose it is possible that non_ep_grads could also be empty, but I can only imagine this happening in extreme cases, so I did not change the non_ep_grads code.

CC @tianyu-l

[tianyu-l]

Thanks, had a minor comment.

[tianyu-l]

This edge case can occur in PP+EP setups when model uses some fully dense and some MoE layers (like DSv3), in which case some pp ranks may not be assigned any MoE layers.

Oh makes sense to me. Could you actually put this example edge case in the comment too? I think it'd be very helpful.

I suppose it is possible that non_ep_grads could also be empty, but I can only imagine this happening in extreme cases, so I did not change the non_ep_grads code.

I think this is not possible if a PP stage always

1. contains any non-MoE params
2. contains full MoE modules -- the shared expert and router.gate will be non_ep_params anyways

[garrett361]

Expanded the comment.

I think this is not possible if a PP stage always [...]

Yeah, I was imagining very extreme cases where PP is very granularly applied an somehow a PP rank only ends up owning MoE layers and nothing else. Can't happen for any model or parallelism you could setup with torchtitan:main today, for sure. I was mostly just explaining why I only touched the ep_grads code.

[garrett361]

Need anything else from me on this one @tianyu-l ?

[tianyu-l]

oh sorry forgot to merge :)

[tianyu-l]

oh sorry forgot to merge :)

[garrett361]

Np, thanks!