TransformerEngine : Add test with FSDP (and updates to ddp_wrapper in test_ddp.py)

[kshitij12345]

This PR adds test for using TE executor in FSDP and verifies it against Eager + TE. Also we update the ddp_wrapper to allow wrapping with different pytest-fixture besides bucket_size_in_mb (which errored when I tried to add a different pytest-fixture).

PR #80 description details of how TE automatically takes care of syncing FP8 meta-data in distributed setting.

Also, I have verified it on a larger model using the available benchmarking script
cmd for benchmark:

torchrun --nproc-per-node=2 thunder/benchmarks/benchmark_litgpt.py --compile thunder+nvfuser+transformerengine+cudnn --n_layers=10 --distributed_mode=fsdp

Numbers are on RTX 6000

Without TE

iter 41: loss 4.6562, iter time: 3180.77ms, t: 4096
iter 42: loss 4.6250, iter time: 3202.35ms, t: 4096
iter 43: loss 4.6562, iter time: 3172.88ms, t: 4096
iter 44: loss 4.6562, iter time: 3181.55ms, t: 4096
Model name: Llama-2-7b-hf
Seq Length: 4096
Micro BS: 1
Global BS: 2
Number of Layers: 10
Number of parameters: 1.14B
Distributed Mode: fsdp
Sharding Mode: zero2
Bucketing: none
Compiler: thunder+nvfuser+cudnn
Average iter time: 3187.63 ms
Memory used: 30.56 GB
Tokens/s: 2570.17
Tokens/s/GPU: 1285.09
TFLOP/s: 38.40

With TE

iter 42: loss 4.6562, iter time: 3025.66ms, t: 4096
iter 43: loss 4.6562, iter time: 3030.40ms, t: 4096
iter 44: loss 4.6562, iter time: 3018.83ms, t: 4096
Model name: Llama-2-7b-hf
Seq Length: 4096
Micro BS: 1
Global BS: 2
Number of Layers: 10
Number of parameters: 1.14B
Distributed Mode: fsdp
Sharding Mode: zero2
Bucketing: none
Compiler: thunder+nvfuser+transformerenginevfuser+cudnn
Average iter time: 3024.72 ms
Memory used: 37.25 GB
Tokens/s: 2708.12
Tokens/s/GPU: 1354.06
TFLOP/s: 40.47

[kshitij12345]

@carmocca please have a look and stamp the PR, thanks!