Implement no_sync for thunder.distributed.fsdp (PR2457)

[crcrpar]

tldr

Enables no_sync for thunder.jit(thunder.distributed.fsdp(model)). The accompanied changes are:

• new argument of return_none_instead_of_grads of ThunderFunction.forward
  • This could be eliminated once a TraceCtx's bound symbols are not deleted even if it just returns one or more Nones
• removal of no_sync check before applying dist_prims.synchronize to args and kwargs
  • FSDP's forward needs this prim for its param AllGather
  • [ddp] visitor_transform removes dist_prims.all_reduce, dist_prims.wait, and preaveraging when no_sync
  • [fsdp] visitor_transform removes comms and puts dist_prims.stash_grad_for_fsdp and optional param AllGather when no_sync
    • The generated trace and its executable python code return unsynchronized unsharded gradients.
    • The prim's implementation accumulates the grads as param._thunder_fsdp_unsharded_grad.
    • ThunderFunction's backward returns Nones instead of such grads to avoid shape mismatch between params and unsharded grads.

---

as of fa61c49

• llama-2-7b-hf
• world size 8 H100s
• micro batch size 1
• global batch size 32
• gradient accumulation 4
• no bucketing (of AllGather and ReduceScatter)

zero2

command: torchrun --nproc-per-node=8 thunder/benchmarks/benchmark_litgpt.py --compile=thunder_inductor --distributed_mode=fsdp --nsys_enabled=False --micro_batch_size=1 --global_batch_size=32 --skip_data_sync <false|true> --model_name=Llama-2-7b-hf --shard_mode=zero2 --bucketing_mode=none --json_path "<filename>.json" --return_metrics_as_json=true

	w/ no_sync	w/o no_sync
tokens/sec	82713.0	80341.0
memory consumption [GB]	65.6	40.3

zero3

command: torchrun --nproc-per-node=8 thunder/benchmarks/benchmark_litgpt.py --compile=thunder_inductor --distributed_mode=fsdp --nsys_enabled=False --micro_batch_size=1 --global_batch_size=32 --skip_data_sync <false|true> --model_name=Llama-2-7b-hf --shard_mode=zero3 --bucketing_mode=none --json_path "<filename>.json" --return_metrics_as_json=true

	w/ no_sync	w/o no_sync
tokens/sec	77839.0	75511.9
memory consumption [GB]	52.5	27.1

[crcrpar]

[RFC] new argument to ThunderFunction.apply

[crcrpar]

cc @parthmannan

[crcrpar]

@t-vi this is ready for merge

[crcrpar]

How is HF's Accelerate or LIghtning's Fabric wrapping PyTorch's no_sync?

• fabric: https://lightning.ai/docs/fabric/stable/advanced/gradient_accumulation.html
• accelerate:
  • no_wync wrapper: https://huggingface.co/docs/accelerate/en/concept_guides/gradient_synchronization
  • accumulate, abstraction of an optimizer step with gradient accumulation: https://huggingface.co/docs/accelerate/v0.29.3/en/package_reference/accelerator#accelerate.Accelerator.accumulate

[t-vi]

Thank you, @crcrpar @IvanYashchuk