compiled rms_norm not numerically accurate (fails to produce good loss curves) when run under tp #497
Description
When doing large scale runs, found that compiled_rmsnorm was producing aberrant loss curves compared to tp or async tp with rmsnorm.
Verified this reproes with small scale and thus opening issue for tracking.
Compiled rmsnorm works fine under 1D FSDP.
Easy repro is to run debug model and 50 steps with compiled rms_norm with tp, vs tp with rms norm:
fsdp + compile + rmsnorm:
fspd + tp (2) + compile + rmsnorm:
Issue shown here:
fsdp + tp(2) + compiled rms norm:
Finally, compiled rmsnorm + fsdp + compile (1D) works fine:
