[not for land] TE experiments, take 2

[vkuzo]

Summary:

Test Plan:

// on a machine with 8 H100 GPUs...
// PT - bf16 + compile
with-proxy CONFIG_FILE="./train_configs/llama3_8b.toml" ./run_llama_train.sh --activation_checkpoint.mode none --training.seq_len 6144 --training.compile
// PT - f8 compute + compile
with-proxy CONFIG_FILE="./train_configs/llama3_8b.toml" ./run_llama_train.sh --activation_checkpoint.mode none --training.seq_len 6144 --training.compile --float8.enable_float8_linear
// PT - f8 compute + comms + compile
with-proxy CONFIG_FILE="./train_configs/llama3_8b.toml" ./run_llama_train.sh --activation_checkpoint.mode none --training.seq_len 6144 --training.compile --float8.enable_float8_linear --float8.enable_fsdp_float8_all_gather --float8.precompute_float8_dynamic_scale_for_fsdp
// TE + compile - doesn't work OOB, needs a workaround
// TE - bf16
with-proxy CONFIG_FILE="./train_configs/llama3_8b.toml" ./run_llama_train.sh --training.use_te --activation_checkpoint.mode none --training.seq_len 6144
// TE - float8
with-proxy CONFIG_FILE="./train_configs/llama3_8b.toml" ./run_llama_train.sh --training.use_te --training.use_te_float8 --activation_checkpoint.mode none --training.seq_len 6144

Reviewers:

Subscribers:

Tasks:

Tags: