FP8BlockQuantizer not work on TE

[fy1214]

Describe the bug

1. Issues Caused by Enabling Tensor Parallelism (TP)
   Phenomenon: When enabling fp8_recipe == blockwise with TP:

• Error: AssertionError: All-gather requires quantizable tensor for quantizer Float8BlockQuantizer
• If this assertion is commented out and SP (Sequence Parallelism) is disabled:
  Error: ValueError: When using expert parallelism and tensor parallelism, sequence parallelism must be used
• If this assertion is commented out and SP is enabled:
  Error: CUDA Error: an illegal memory access was encountered
  Temporary Solution: Disable TP for MOE training.

2. Issues Caused by Enabling CPU Offload
   If optimizer_cpu_offload is enabled, it requires:

• Not enabling --fp8-param-gather
• And args.fp8_recipe == "delayed"
  Problem: If --fp8-param-gather is not enabled, FP8 training ends up requiring higher GPU memory than BF16 training.

Environment details

If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:

• Pytorch version 2.8.0+cu129
• Python version 3.12
• Transformer Engine version 2.8.0
• CUDA version 12.9

[guyueh1]

@fy1214 About the error All-gather requires quantizable tensor for quantizer Float8BlockQuantizer, the blockwise quantizer requires the tensor shape [M, N] to meet N%128==0 and M%128==0, the condition is defined here. In that specific place where assertion is through, it's trying to determine of input to LayernormLinear is quantizable. The input is expect to have shape [S/CP/TP, H] where H is the hidden_size of model, S is the pretrain sequence length per sample; if you are using sequence packing, it's more complicated, but the key is to make sure your packed sequence, after sharded to CP and TP ranks, is divisible by 128. That needs to be provided by your framework, can you check if your framework already guarantee that?

About ValueError: When using expert parallelism and tensor parallelism, sequence parallelism must be used, this is not a bug, it's a requirement, TE doesn't support EP+TP without SP. Please make sure you turn on SP when you have TP in attention and EP in mlp.

I will follow up on testing optimizer cpu offloading later.

[guyueh1]

Sorry, my bad, for the first error, the first dimension doesn't need to be divisible by 128 because it's per-channel quantized. Let me look into this.

[fy1214]

@fy1214 About the error All-gather requires quantizable tensor for quantizer Float8BlockQuantizer, the blockwise quantizer requires the tensor shape [M, N] to meet N%128==0 and M%128==0, the condition is defined here. In that specific place where assertion is through, it's trying to determine of input to LayernormLinear is quantizable. The input is expect to have shape [S/CP/TP, H] where H is the hidden_size of model, S is the pretrain sequence length per sample; if you are using sequence packing, it's more complicated, but the key is to make sure your packed sequence, after sharded to CP and TP ranks, is divisible by 128. That needs to be provided by your framework, can you check if your framework already guarantee that?

About ValueError: When using expert parallelism and tensor parallelism, sequence parallelism must be used, this is not a bug, it's a requirement, TE doesn't support EP+TP without SP. Please make sure you turn on SP when you have TP in attention and EP in mlp.

I will follow up on testing optimizer cpu offloading later.

Thanks.

[fy1214]

@fy1214 About the error All-gather requires quantizable tensor for quantizer Float8BlockQuantizer, the blockwise quantizer requires the tensor shape [M, N] to meet N%128==0 and M%128==0, the condition is defined here. In that specific place where assertion is through, it's trying to determine of input to LayernormLinear is quantizable. The input is expect to have shape [S/CP/TP, H] where H is the hidden_size of model, S is the pretrain sequence length per sample; if you are using sequence packing, it's more complicated, but the key is to make sure your packed sequence, after sharded to CP and TP ranks, is divisible by 128. That needs to be provided by your framework, can you check if your framework already guarantee that?

About ValueError: When using expert parallelism and tensor parallelism, sequence parallelism must be used, this is not a bug, it's a requirement, TE doesn't support EP+TP without SP. Please make sure you turn on SP when you have TP in attention and EP in mlp.

I will follow up on testing optimizer cpu offloading later.

The first error we fix it by padding our input to 128. Now only the optimizer_cpu_offload problem left.

[guyueh1]

@fy1214 Padding is the correct way, I confirmed with TransformerEngine dev team that it is a requirement that seqlen/TP is multiple of block size, this is to guarantee that quantization is correctly performed for backward.

[guyueh1]

@fy1214 I tested offloading with (block-wise recipe + fp8_param_gather=true), when use_precision_aware_optimizer=False it works fine for me. When use_precision_aware_optimizer=True I get

rank1]:   File "/opt/megatron-lm/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py", line 57, in __init__
[rank1]:     super(HybridDeviceOptimizer, self).__init__(
[rank1]:   File "/usr/local/lib/python3.12/dist-packages/torch/optim/optimizer.py", line 400, in __init__
[rank1]:     self.add_param_group(cast(dict, param_group))
[rank1]:   File "/usr/local/lib/python3.12/dist-packages/torch/_compile.py", line 53, in inner
[rank1]:     return disable_fn(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 1005, in _fn
[rank1]:     return fn(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/usr/local/lib/python3.12/dist-packages/torch/optim/optimizer.py", line 1112, in add_param_group
[rank1]:     raise TypeError(
[rank1]: TypeError: optimizer can only optimize Tensors, but one of the params is NoneType

Is this matching your observation?

[guyueh1]

Nevermind, according to https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/training/arguments.py#L1165-L1173 your use case is unsupported by megatron. Specifically, when using CPU offloading optimizer, you must turn on use_precision_aware_optimizer, and disable fp8_param_gather unless it's delayed scaling. The implementation for Fp8 parameter optimizer for other scaling recipes is missing in TransformerEngine.

[tkob-vh]

@fy1214 Padding is the correct way, I confirmed with TransformerEngine dev team that it is a requirement that seqlen/TP is multiple of block size, this is to guarantee that quantization is correctly performed for backward.

@guyueh1 Hello. I found that the assert_dim_for_all_gather is included in pr 2177, and merged in release 2.8. Before release 2.8, it is not needed to assert the first dim should also be divided by 128. Does it means that before release 2.8 it is not correct?
I'm using fp8 blockwise quantization on dense models, enabling tensor parallelism and sequence parallelism. I commented out this assertion and it seems that it works fine for me.🤔
cc @ksivaman

[guyueh1]

@tkob-vh according to @zhongbozhu , before release 2.8 the assertion was not included but the behavior was not correct either for cases when SL/TP is not divisible by 128. In other words, TE never properly supported activation block scaling when a block is splitted to different TP ranks, because that would need to collect the block amax cross GPU and we don't have such utility. Before 2.8 it was just silently producing wrong results.

[zhongbozhu]

This constraint is only there because the TP AG communication for FP8 recipe will be in FP8 for better performance, half the network communication volume.

The numerical skew here should be small, but non-zero. Image a local tensor shape [64, 128]. The vanilla implementation is that we do AG in BF16 to get a tensor [128, 128], and then quantize locally. For forward pass, we are quantizing 1x128, but for backward pass it's 128x1 (equivalent to transpose the high precision input and then perform 1x128 scaling, but high precision tensor transpose is super expensive, so we do 128x1 block scaling and transpose in shared memory instead).

transpose ( block_scaling_128x1 ( x_bf16) ) == block_scaling_1x128( transpose ( x_bf16 ) )

If we want to quantize before all-gather. For backward pass, we are performing the 128x1 scaling on the 64 axis. To be strictly correct in TP AG, we have to find local max vertically along 64 dimension, do amax reduction among TP group, and then quantize locally to FP8, and then finally we can do AG in FP8.

Without this assert, we are essentially performing 64x1 block scaling for backward, the missing 64 elements will be zero filled by the quantization kernel automatically, this won't be exactly the same with 128x1 block scaling when two TP ranks have different max values every 64 elements. But the strictly correct way of doing it is the way I mentioned above, which is a lot of steps, and will be pretty slow (fragmented small kernels, CPU overhead, etc.).

We could, in theory, add back the BF16 AG fall back at the cost of some performance penalty & more if branches which will make the code look messy :) Deepseek FP8 block scaling recipe is a special case where the block size 128 is actually pretty large and handling corner cases is trickier here like this issue described.

That being said, I also think keeping this error message without silently falling back to the slow BF16 AG path is maybe also not a bad idea to remind users about using a more aligned shape. Since TP AG is only applied to dense linear layers, in which enforcing a regular shape shouldn't be an issue right? Can you remind me of any cases where it's difficult to enforce a more aligned shape in the token dimension M?