Fix PP clip_grad_norm #649
Conversation
facebook-github-bot
added
the
CLA Signed
torchtitan/clip_grad_nrom.py
Outdated
| total_norm **= norm_type | ||
| dist.all_reduce(total_norm, op=dist.ReduceOp.SUM, group=pp_mesh.get_group()) | ||
| total_norm **= 1.0 / norm_type |
There was a problem hiding this comment.
This is only correct for p-norm where p is int/float (e.g. this does not work for inf-norm).
Intuitively, we should somehow allow us to change total_norm to have _NormPartial(norm_type) placement over the PP mesh, and we should redistribute it to Replicate() to reuse DTensor's own logic. cc: @tianyu-l @wz337 @XilunWu is there any way to do this?
There was a problem hiding this comment.
@awgu in torchtitan, PP is done by first manually splitting the Llama model then wrap in PP wrapper. If my understanding is correct, none of them converts models weights into DTensor with PP mesh placement
There was a problem hiding this comment.
For now, will just add a if-statement to check for inf norm
There was a problem hiding this comment.
PP is done by first manually splitting the Llama model then wrap in PP wrapper. If my understanding is correct, none of them converts models weights into DTensor with PP mesh placement
this is true, however after splitting the llama model we do also apply SPMD parallelisms (TP and DP) which will use DTensor for parameters. So by the time we are clipping norms, parameters would be DTensors. (if we use PP only and no TP/DP, then there would be no DTensor parameters)
There was a problem hiding this comment.
@wconstab You are right about SPMD (TP and/or DP) but the DTensor in that case has placement on TP/DP mesh not on PP mesh. We still need to do some form of reduce either explicitly (same as my code) or wrap the total_norm in DTensor on PP mesh (maybe via DTensor.redistribute?)
|
|
||
| if pp_mesh is not None: | ||
| if isinstance(total_norm, DTensor): | ||
| # will reach here if PP + other parallelism is used. If only using PP, total_norm will be a local tensor |
There was a problem hiding this comment.
@wconstab the if statement are now separated. Please let me know if it make sense to you
torchtitan/utils.py
Outdated
| # if total_norm is a DTensor, the placements must be `torch.distributed._tensor.ops.math_ops._NormPartial` | ||
| # we can simply reduce the DTensor to get the total norm in this tensor's process group | ||
| # and then convert it to a local tensor | ||
| total_norm = total_norm.redistribute( |
There was a problem hiding this comment.
nit: use total_norm.full_tensor() instead?
There was a problem hiding this comment.
@H-Huang Good suggestion. Updated!
@awgu please let me know if the update make sense and feel free to resolve the change request if you believe it is good enough |
This PR fixes the pipeline parallel incomplete gradient norm issue. Refer to pytorch#596 for more details.
This PR fixes the pipeline parallel incomplete gradient norm issue. Refer to #596 for more details.