enable FSDP2 + fp8 all-gather and fix TP fp8 all-gather

[weifengpy]

we have landed fp8 all-gather optimizations in float8_experimental pytorch-labs/float8_experimental#266

this PR proposes torchtitan changes. also include fp8 in CI

from float8_experimental.fsdp_utils import precompute_float8_dynamic_scale_for_fsdp
# inside the training loop
model(input).sum().backward()
optim.step()
precompute_float8_dynamic_scale_for_fsdp(model)

FSDP2 fp8 all-gather are added to CI

CONFIG_FILE="./train_configs/debug_model.toml" ./run_llama_train.sh --training.enable_fp8_linear
CONFIG_FILE="./train_configs/debug_model.toml" ./run_llama_train.sh --training.enable_fp8_linear --training.enable_fsdp_fp8_all_gather
CONFIG_FILE="./train_configs/debug_model.toml" ./run_llama_train.sh --training.enable_fp8_linear --training.enable_fsdp_fp8_all_gather --training.precompute_float8_dynamic_scale_for_fsdp

TP fp8 all-gather are locally tested. will add them to CI after uploading a new tokenizer with vacab size 2560 (divisible by 16)

CONFIG_FILE="./train_configs/llama3_8b.toml" NGPU=4 ./run_llama_train.sh --training.enable_fp8_linear --training.data_parallel_degree 1 --training.tensor_parallel_degree 4
CONFIG_FILE="./train_configs/llama3_8b.toml" NGPU=4 ./run_llama_train.sh --training.enable_fp8_linear --training.data_parallel_degree 2 --training.tensor_parallel_degree 2

precompute scales after optimizer.step

FSDP2 pre-all-gather do not have any small all-reduces

TODO

• upload tokenizer with vacab size 2560 to enable CI on TP fp8 all-gather
• torch.compile complains about fp8
• add delayed scaling and brainstorm about best config option to express fp8
• compare perf between delayed scaling and dynamic scaling https://github.com/pytorch-labs/float8_experimental/pull/312/files

[weifengpy]

added followings to CI

• 1D fsdp original dtype all-gather
• 1D fsdp fp8 all-gather
• 1D fsdp fp8 all-gather with precomputed dynamic scales

need follow ups to enable TP fp8 all-gather in CI: current CI tokenizer has 2556, not divisible by 16) #461

• 1D TP fp8 all-gather
• 2D FSDP + TP fp8 all-gather

[tianyu-l]

maybe noob question: does all-gather always happen in bf16, or is it determined by param_dtype of FSDP2

[weifengpy]

good question! I should change this to "all-gather in original dtype"

when mixed_precision is turned off (no param_dtype configed), FSDP2 communicate in model's original dtype

when mixed_precision is turned on (param_dtype=xxx), FSDP2 communicates according to param_dtype

[tianyu-l]

maybe a noob question: could you briefly explain what this is doing?
I wonder since we are already using context functions for FP8, can we have a context and run it in a .step() function here, just like optimizer, lr scheduler, and profiler. This would make the code consistent.

[weifengpy]

could you briefly explain what this is doing

precompute_float8_dynamic_scale_for_fsdp is a for-loop over model.parameters(). it issues a single all-reduce for all parameters, ie abs(max(param)) for param in model.parameters() and save amax/scale as param._precomputed_scale. this speed up the training loop since we do not need to compute amax/scale for each parameters in the training loop

we are already using context functions for FP8

do you refer to set_enable_fsdp_fp8_all_gather ? That's for model intiaitialization where we swap nn.Linear with user-defined float8 linear. precompute_float8_dynamic_scale_for_fsdp is for training loop

[weifengpy]

per suggestion, raise error if use_fp8_linear=False or enable_fsdp_fp8_all_gather =False

[yifuwang]

noob q: do we eventually want to just put this in fsdp2?

[awgu]

It has to be done after optimizer step (since parameter values change). Are you suggesting to run this in the root module's pre-forward?

[yifuwang]

Yeah anywhere between the n-1th optimizer step and the first all-gather in the nth step where fsdp2 has control (if there's any).

[awgu]

That makes sense. I think one concern is that FSDP is agnostic to the fp8 all-gather. FSDP does not know that the fsdp_pre_all_gather and fsdp_post_all_gather of the Float8Linear.weights are implemented to do fp8 all-gather, so at best, the user still would need to register a module forward pre-hook or something to run this method.

[yifuwang]

Ah I see. Somehow I thought fsdp2 was fp8-aware

[weifengpy]

add comment to explain precompute_float8_dynamic_scale_for_fsdp

[weifengpy]

done

[tianyu-l]

Left some comments to record what we discussed offline.

[tianyu-l]

As discussed offline, let's refactor fp8 configs, e.g. have a dedicated field for enabling fp8 or not.

[weifengpy]

renamed fp8_linear to enable_fp8_linear

[drisspg]

I think one thing to note is that right now this is a boolean which will swap to the default float8 recipe
Dynamic scaling x Tensor wise ScalingGranularity x all tensors involved in the matmul [ input, weight, grad]

I think we should brainstorm on an elegant solutions for users to express their desired config here

[weifengpy]

good question. evetually we might have to expose args/kwargs from swap_linear_with_float8_linear for flexibility

[tianyu-l]

discussed offline: please check if it makes sense to enable it only when dp_degree > 1

[weifengpy]

added check on parallel_dims.dp_enabled

[tianyu-l]

can remove this to favor simplicity if it is a no-op flag when fp8_linear=False

[weifengpy]

removed ValueError on enable_fp8_linear=False

[tianyu-l]

discussed offline: can refactor to make it simpler

[weifengpy]

removed ValueError when enable_fp8_linear/enable_fsdp_fp8_all_gather=False

[weifengpy]

fp8_linear == "dynamic" is outdated after recent unification of dynamic/delayed scaling (Float8Linear)
#436

update it in this PR to make TP fp8 all-gather work again

EDIT: Will enable TP in CI to prevention after having a new tokenizer with vacab size 2560

[drisspg]

noop Q: do we need this in a context manager to make testing + resetting easier?

[weifengpy]

hmm. set_enable_fsdp_fp8_all_gather is a context manager right now. do you mean "why" it should be a context manager ?

EDIT: I also see you mentioned "make testing + resetting easier", which answered why. so I am not sure if it's a question for me

[drisspg]

Should be "only supports dynamic scaling" right?

[weifengpy]

good catch! will change it

[tianyu-l]

nit: comment for 2D

[weifengpy]

in the end I have to remove 2D from this PR. current CI tokenizer has vacab size = 2556. However, fp8 gemm need the vacab size to be divisible by 16 #461

I can follow up with you on how to have a tokenizer with vacab size = 2560 to unblock 1D TP + fp8, and 2D + fp8 in CI

[tianyu-l]

The PR looks good to me, except that there is an ongoing discussion on what's the best way to support fp8 GEMM on tensors of irregular shapes (when not divisible by 16) e.g. coming from TP + test tokenizer model in this repo (see #461). Maybe we should allow irregular shapes but with potential perf regression. Will need to discuss more on this later.

[tianyu-l]

let's remove this redundant check

[wanchaol]

@weifengpy I think we should hide this import to the path where enable_fp8_allgather path happened?

The problem here is that for every feature that requires an additional install from other dependency, we should try to hide the import to the path that uses it instead of import it globally, otherwise for users who didn't install the float8_experimental, if they rebase, and it would just fail to train for them.

Please submit a follow up PR to fix this

[weifengpy]

got you. I am moving it from top-level to if-else now #464

thanks for the timely reminder