TransformerEngine - Intermediate tensor sharding

[kshitij12345]

TransformerEngine added the ability to shard intermediate activation tensors in v1.8. Currently, we save global/world sized activation for backward pass. Using this, we can lower the peak memory usage at the cost of added comms - as we will shard these intermediate tensor and gather them before the backward computation.

TE PR: NVIDIA/TransformerEngine#687

In this PR, we use make this option opt-in using thunder.jit compile argument - fp8_shard_intermediate_activation.

Example usage: model = thunder.jit(model, executors=executors, fp8_shard_intermediate_activation=True)

Testing

Updated the distributed test to use this option. Have tested with existing tests in test_transformer_engine_executor.py and test_ddp.py -k transformer with TE v1.7 (current stable), v1.8 and v1.9 (current main).

Benchmark
Command -

torchrun --nproc_per_node=8 --nnodes=1 thunder/benchmarks/benchmark_litgpt.py --return_metrics_as_json=True --json_path=/tmp/benchmark_litgpt_data.json --distributed_mode=fsdp --shard_mode=zero3 --model_name=Llama-2-7b-hf --micro_batch_size=1 --compile=thunder_inductor_cat_transformerengine_cudnn --nsys_enabled=False --dynamic=False

Without FP8 Intermediate Sharding

Average iter time: 282.47 ms
Memory used: 52.92 GB

With FP8 Intermediate Sharding

Average iter time: 341.67 ms
Memory used: 44.05 GB

Patch to enable sharding in `benchmark_litgpt.py`

diff --git a/thunder/benchmarks/benchmark_litgpt.py b/thunder/benchmarks/benchmark_litgpt.py
index bad6ef74..947295fc 100644
--- a/thunder/benchmarks/benchmark_litgpt.py
+++ b/thunder/benchmarks/benchmark_litgpt.py
@@ -341,7 +341,7 @@ class Benchmark_litGPT:
 
                 executors.insert(0, transformer_engine_ex)
 
-            model = thunder.jit(model, executors=executors)
+            model = thunder.jit(model, executors=executors, fp8_shard_intermediate_activation=True)
 
         elif self.compile != "eager":
             raise ValueError(f"Invalid compile option: {self.compile}")

[t-vi]

Thank you @kshitij12345 @IvanYashchuk