Add numerics comparison script

[xmfan]

Stacked PRs:

• ->Add numerics comparison script #259

---

Add numerics comparison script from Devmate

Intended usage: python examples/run_ds3_numerics_check.py

Which creates a temp folder e.g. /tmp/ds3_numerics_check_ufc54rz8, then runs:

torchrun --standalone --nproc-per-node 4 examples/example_ds3_local_map.py --rng-seed 42
torchrun --standalone --nproc-per-node 8 examples/example_ds3_pp.py --rng-seed 42
diff out/0/weights.log out/1/pp_weights.log
diff out/0/diff.log out/1/diff.log

It currently warns, because there's a silly mismatch wrt buffer registration, where we have 4 freq_cis in the pp model. Once fixed, I can change it back to a RuntimeError, and we can add it into CI.

[sanketpurandare]

Have you verified numerics with the flag --use-loss-fn and by changing the schedules using --schedule-name?

[xmfan]

nope, for loss_fn, we would need to add it to the non-pp script

and for schedule names, i'll need to validate them. will add a CLI flag to switch between schedules for now

[xmfan]

@sanketpurandare interleave1f1b and zbpp are both fine, dualpipev needs some minor logging changes. will do separately and add into ci for next pr