Add option for run-to-run deterministic and add optional numerics logging

[xmfan]

Stacked PRs:

• ->Add option for run-to-run deterministic and add optional numerics logging #235

---

Add option for run-to-run deterministic and add optional numerics logging

Example log:
tlp torchrun --standalone --nproc-per-node 8 examples/example_ds3_pp.py --rng-seed 1234 | pastry; tlp torchrun --standalone --nproc-per-node 8 examples/example_ds3_pp.py --rng-seed 1234 | pastry

https://www.internalfb.com/intern/diffing/?paste_number=2027323783

[sanketpurandare]

Maybe we should add this to all the graph module calls

[xmfan]

you mean the backward? I didn't add it since I couldn't test it on the base commit

[sanketpurandare]

Yeah, once I land #237 we can add it for full_bw, bw_dI, bw_dW, unshard and reduce_grad.

[sanketpurandare]

Let's say we have 8 ranks in total, they will all initialize their modules. Since each rank initializes a different part of the model, it is hard to compare it with a single rank implementation for numerics debugging. We should have a solution similar to what @wconstab used in torchtitan. Creating a seed checkpoint and using that for PP runs.

[xmfan]

By seed checkpoint, do you mean saving and loading random weights generated from a rng seed? I was thinking of just resetting the seed for weights init

[sanketpurandare]

So if pp has 8 stages, you would do init_weights for each one of them using the same seed? My concern is how would you compare the pp_runtime with spmd only for numerics?

[xmfan]

If we cut our stages at nm module boundaries, and init weights in the same order, we could reset the seeds at the same cuts during the spmd init weights.

For supporting arbitrary stage splits, I would need to know more about how we would implement their init_weights and checkpointing. So I put that aside for now.

[sanketpurandare]

Could you also add an example that saves the params after init and the grads after after accumulating grads by running microbatches in spmd? Analogously, pp also saves the params after init and grads after running the step and finally a script that compares both?

[xmfan]

Yup, I'm changing up example_ds3_local_map.py to use real tensors to be the SPMD microbatch + accumulate grad steps baseline. And I have a script to diff the outputs of DebugInterpreter that I was thinking of landing separately from this PR.

[sanketpurandare]

Just land after #237 and add the DebugInterpreter to other graph_module calls if required.