Hopper + blockwise fp8 + tp=2 + ep dispatch=alltoall got RuntimeError: GEMM with Float8BlockwiseQTensor requires GEMM_READY format
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
**Describe the bug**
Got runtime error when running moe megatron training on Hopper with:
--tensor-model-parallel-size 2
--expert-model-parallel-size 2
--expert-tensor-parallel-size 1
--moe-token-dispatcher-type alltoall
--fp8-param-gather
--fp8-recipe blockwise
--fp8-format e4m3
[rank4]: File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1767, in _wrapped_call_impl
[rank4]: return self._call_impl(*args, **kwargs)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _call_impl
[rank4]: return forward_call(*args, **kwargs)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/dnn_training_sys/haitao/github/Megatron-LM/megatron/core/transformer/transformer_layer.py", line 435, in forward
[rank4]: output = self._forward_mlp(hidden_states, kwargs.get("inference_context", None))
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/dnn_training_sys/haitao/github/Megatron-LM/megatron/core/transformer/transformer_layer.py", line 618, in _forward_mlp
[rank4]: mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1767, in _wrapped_call_impl
[rank4]: return self._call_impl(*args, **kwargs)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _call_impl
[rank4]: return forward_call(*args, **kwargs)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/dnn_training_sys/haitao/github/Megatron-LM/megatron/core/transformer/moe/moe_layer.py", line 289, in forward
[rank4]: output, mlp_bias = custom_forward(hidden_states)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/dnn_training_sys/haitao/github/Megatron-LM/megatron/core/transformer/moe/moe_layer.py", line 271, in custom_forward
[rank4]: output, shared_expert_output, mlp_bias = self.experts_compute(
[rank4]: ^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/dnn_training_sys/haitao/github/Megatron-LM/megatron/core/transformer/moe/moe_layer.py", line 225, in experts_compute
[rank4]: self.token_dispatcher.dispatch_postprocess(hidden_states, probs)
[rank4]: File "/dnn_training_sys/haitao/github/Megatron-LM/megatron/core/transformer/moe/token_dispatcher.py", line 651, in dispatch_postprocess
[rank4]: self.shared_experts.linear_fc1_forward_and_act(global_input_tokens)
[rank4]: File "/dnn_training_sys/haitao/github/Megatron-LM/megatron/core/transformer/moe/shared_experts.py", line 179, in linear_fc1_forward_and_act
[rank4]: intermediate_parallel, bias_parallel = self.linear_fc1(self.cached_fc1_input)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1767, in _wrapped_call_impl
[rank4]: return self._call_impl(*args, **kwargs)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _call_impl
[rank4]: return forward_call(*args, **kwargs)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/dnn_training_sys/haitao/github/Megatron-LM/megatron/core/extensions/transformer_engine.py", line 426, in forward
[rank4]: out = super().forward(x, is_first_microbatch=_is_first_microbatch)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 893, in _fn
[rank4]: return fn(*args, **kwargs)
[rank4]: ^^^^^^^^^^^^^^^^^^^
[rank4]: File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/module/linear.py", line 1494, in forward
[rank4]: out = linear_fn(*args)
[rank4]: ^^^^^^^^^^^^^^^^
[rank4]: File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/module/linear.py", line 309, in forward
[rank4]: gemm_out, *_, reduce_scatter_out = general_gemm(
[rank4]: ^^^^^^^^^^^^^
[rank4]: File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/cpp_extensions/gemm.py", line 120, in general_gemm
[rank4]: raise RuntimeError("GEMM with Float8BlockwiseQTensor requires GEMM_READY format")
[rank4]: RuntimeError: GEMM with Float8BlockwiseQTensor requires GEMM_READY format
If use flex for dispatch, no error.
**Steps/Code to reproduce bug**
Please list *minimal* steps or code snippet for us to be able to reproduce the bug.
A helpful guide on on how to craft a minimal bug report http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports.
**Expected behavior**
No error.
**Additional context**
Add any other context about the problem here.
Contributor guide
Assessment
This issue has not been assessed yet.