From 8634f53507571b69ba1762d8e25bef01e40d36cd Mon Sep 17 00:00:00 2001 From: kunlunl Date: Sat, 11 Oct 2025 21:23:40 +0800 Subject: [PATCH 1/9] Add support for FP8 param + CUDA Graph --- .../distributed/distributed_data_parallel.py | 50 +++++++++++-------- .../core/distributed/param_and_grad_buffer.py | 20 ++++++-- megatron/core/fp8_utils.py | 9 ++++ 3 files changed, 53 insertions(+), 26 deletions(-) diff --git a/megatron/core/distributed/distributed_data_parallel.py b/megatron/core/distributed/distributed_data_parallel.py index df5bccd71ca..df1d7ae94db 100644 --- a/megatron/core/distributed/distributed_data_parallel.py +++ b/megatron/core/distributed/distributed_data_parallel.py @@ -8,7 +8,7 @@ from .. import parallel_state from ..config_logger import has_config_logger_enabled, log_config_to_disk -from ..fp8_utils import is_float8tensor +from ..fp8_utils import is_float8tensor, post_all_gather_processing from ..process_groups_config import ProcessGroupCollection from ..transformer.cuda_graphs import is_graph_capturing from ..transformer.transformer_config import TransformerConfig @@ -500,26 +500,34 @@ def start_param_sync(self, *unused, force_sync: bool = False, force_dispatch: bo for bucket_group in self.bucket_groups + self.expert_parallel_bucket_groups: bucket_group.start_param_sync(force_sync=force_sync) - # For MXFP8 params, we need to copy the all-gathered param data from the buffer to - # the param.data, since param buffer is not mapped to model params for MXFP8 case. - # The paramaters are cast from bf16 to MXFP8 during copy. - # In the case of "overlap_param_gather=True", the param copy is done - # in "finish_param_sync" stage after zeroing the shared gardient buffers. - if ( - self.ddp_config.reuse_grad_buf_for_mxfp8_param_ag - and not self.ddp_config.overlap_param_gather - ): - for bucket in bucket_group.buckets: - for param in bucket.params: - param_start, param_end = bucket.param_to_index[param] - param_slice = bucket.param_data.view(-1)[param_start:param_end] - param.data.copy_(param_slice.view(param.data.shape)) - # All-gathered params are not needed after being copied to param.data. - # Zero out the param buffer (shared with grad buffer) for gradient accumulation. - # We cannot zero out the entire grad buffer because one grad buffer may - # correspond to multiple param buffers. If we zero out the entire grad buffer, - # it would clear the data of those param buffers that have not yet completed AG. - bucket.param_data.zero_() + + if not self.ddp_config.overlap_param_gather: + # For MXFP8 params, we need to copy the all-gathered param data from the buffer to + # the param.data, since param buffer is not mapped to model params for MXFP8 case. + # The paramaters are cast from bf16 to MXFP8 during copy. + # In the case of "overlap_param_gather=True", the param copy is done + # in "finish_param_sync" stage after zeroing the shared gardient buffers. + if self.ddp_config.reuse_grad_buf_for_mxfp8_param_ag: + for bucket in bucket_group.buckets: + for param in bucket.params: + param_start, param_end = bucket.param_to_index[param] + param_slice = bucket.param_data.view(-1)[param_start:param_end] + param.data.copy_(param_slice.view(param.data.shape)) + # All-gathered params are not needed after being copied to param.data. + # Zero out the param buffer (shared with grad buffer) for gradient + # accumulation. We cannot zero out the entire grad buffer because one grad + # buffer may correspond to multiple param buffers. If we zero out the entire + # grad buffer, it would clear the data of those param buffers that have not + # yet completed AG. + bucket.param_data.zero_() + else: + fp8_params = [] + for bucket in bucket_group.buckets: + for param in bucket.params: + if is_float8tensor(param): + fp8_params.append(param) + if len(fp8_params) > 0: + post_all_gather_processing(fp8_params) def start_grad_sync(self, *unused): """ diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py index 806defa5b34..d34fdebaf75 100644 --- a/megatron/core/distributed/param_and_grad_buffer.py +++ b/megatron/core/distributed/param_and_grad_buffer.py @@ -17,7 +17,12 @@ from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.rerun_state_machine import get_rerun_state_machine -from ..fp8_utils import is_float8tensor, is_mxfp8tensor, modify_underlying_storage +from ..fp8_utils import ( + is_float8tensor, + is_mxfp8tensor, + modify_underlying_storage, + post_all_gather_processing, +) from ..utils import is_torch_min_version, log_on_each_pipeline_stage from .distributed_data_parallel_config import DistributedDataParallelConfig from .reduce_scatter_with_fp32_accumulation import reduce_scatter_with_fp32_accumulation @@ -311,10 +316,7 @@ def finish_param_sync(self, skip_next_bucket_dispatch: bool = False): # For the mxfp8_param with "reuse_grad_buf_for_mxfp8_param_ag=True", # we need to copy the param_data from the shared_param/grad_buffer to param.data # after the param all-gather. - if ( - self.ddp_config.reuse_grad_buf_for_mxfp8_param_ag - and self.ddp_config.overlap_param_gather - ): + if self.ddp_config.reuse_grad_buf_for_mxfp8_param_ag: for bucket in self.buckets: for param in bucket.params: param_start, param_end = bucket.param_to_index[param] @@ -326,6 +328,14 @@ def finish_param_sync(self, skip_next_bucket_dispatch: bool = False): # correspond to multiple param buffers. If we zero out the entire grad buffer, # it would clear the data of those param buffers that have not yet completed AG. bucket.param_data.zero_() + else: + fp8_params = [] + for bucket in self.buckets: + for param in bucket.params: + if is_float8tensor(param): + fp8_params.append(param) + if len(fp8_params) > 0: + post_all_gather_processing(fp8_params) def start_grad_sync(self): """ diff --git a/megatron/core/fp8_utils.py b/megatron/core/fp8_utils.py index c6ea15776bd..18f0187953a 100644 --- a/megatron/core/fp8_utils.py +++ b/megatron/core/fp8_utils.py @@ -406,6 +406,15 @@ def correct_amax_history_if_needed(model: List[torch.nn.Module]): _correct_amax_history_if_needed_impl(model) +def post_all_gather_processing(model_params): + try: + from transformer_engine.pytorch.tensor.utils import post_all_gather_processing + + post_all_gather_processing(model_params) + except ImportError: + pass + + def is_first_last_bf16_layer(config: TransformerConfig, layer_no: int): """Check if the layer is in bf16.""" num_bf16_layers_at_start = ( From 9915b1f463c413bd202df6901c18fd9e9b83dbd8 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Mon, 13 Oct 2025 15:16:51 +0800 Subject: [PATCH 2/9] Update unit test for fp8 params + cuda graph --- tests/unit_tests/test_fp8_param.py | 86 +++++++++++++++++++++++++++++- 1 file changed, 84 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/test_fp8_param.py b/tests/unit_tests/test_fp8_param.py index 0cbba273507..15cc17b89a9 100644 --- a/tests/unit_tests/test_fp8_param.py +++ b/tests/unit_tests/test_fp8_param.py @@ -29,6 +29,15 @@ _SEED = 1234 fp8_available, reason_for_no_fp8 = check_fp8_support() +cuda_graph_supported = False +reason_for_no_cuda_graph = "" +try: + from transformer_engine.pytorch.tensor.utils import post_all_gather_processing + + cuda_graph_supported = True +except ImportError: + reason_for_no_cuda_graph = "Need newer TransformerEngine" + class TestFP8Param: @@ -68,7 +77,15 @@ def model_provider( ) def create_test_args( - self, tp, recipe, sequence_length, micro_batch_size, inference, fp8_param_gather, **kwargs + self, + tp, + recipe, + sequence_length, + micro_batch_size, + inference, + fp8_param_gather, + use_cuda_graph, + **kwargs, ): destroy_global_vars() destroy_num_microbatches_calculator() @@ -102,6 +119,11 @@ def create_test_args( if recipe == "mxfp8" and fp8_param_gather: args.reuse_grad_buf_for_mxfp8_param_ag = True + if use_cuda_graph: + args.external_cuda_graph = True + args.cuda_graph_scope = "attn" + args.cuda_graph_warmup_steps = 0 + for key, value in kwargs.items(): assert hasattr(args, key) setattr(args, key, value) @@ -122,7 +144,13 @@ def get_batch(self, seq_length, micro_batch_size): return input_ids, labels, position_ids, attention_mask, loss_mask def _run_test_helper( - self, tp_size, recipe, inference: bool = False, fp8_param_gather: bool = True, **kwargs + self, + tp_size, + recipe, + inference: bool = False, + fp8_param_gather: bool = True, + use_cuda_graph: bool = False, + **kwargs, ): """Test fp8_param with gpt_model.""" args = self.create_test_args( @@ -132,6 +160,7 @@ def _run_test_helper( self.micro_batch_size, inference, fp8_param_gather, + use_cuda_graph, **kwargs, ) @@ -239,6 +268,15 @@ def run_test(self, tp_size, recipe, inference: bool = False, **kwargs): ) torch.testing.assert_close(loss_list, loss_list_ref, atol=1e-4, rtol=1e-4) + def run_test_with_cuda_graph(self, tp_size, recipe, **kwargs): + loss = self._run_test_helper( + tp_size, recipe, fp8_param_gather=True, use_cuda_graph=True, **kwargs + ) + loss_ref = self._run_test_helper( + tp_size, recipe, fp8_param_gather=True, use_cuda_graph=False, **kwargs + ) + torch.testing.assert_close(loss, loss_ref, atol=0, rtol=0) + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) @pytest.mark.parametrize("tp_size", [2]) @pytest.mark.parametrize("dp_overlap", [(True, True)]) @@ -246,6 +284,14 @@ def test_delayed_scaling(self, tp_size, dp_overlap): kwargs = {"overlap_param_gather": dp_overlap[0], "overlap_grad_reduce": dp_overlap[1]} self.run_test(tp_size=tp_size, recipe="delayed", **kwargs) + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + @pytest.mark.parametrize("tp_size", [2]) + @pytest.mark.parametrize("dp_overlap", [(True, True)]) + @pytest.mark.skipif(not cuda_graph_supported, reason=reason_for_no_cuda_graph) + def test_delayed_scaling_with_cuda_graph(self, tp_size, dp_overlap): + kwargs = {"overlap_param_gather": dp_overlap[0], "overlap_grad_reduce": dp_overlap[1]} + self.run_test_with_cuda_graph(tp_size, "delayed", **kwargs) + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) @pytest.mark.skipif(not is_te_min_version("2.2.0"), reason="TE 2.2.0 is required") @pytest.mark.parametrize("tp_size", [2]) @@ -260,6 +306,15 @@ def test_tensorwise_scaling(self, tp_size, dp_overlap): def test_tensorwise_scaling_inference(self, tp_size): self.run_test(tp_size=tp_size, recipe="tensorwise", inference=True) + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + @pytest.mark.skipif(not is_te_min_version("2.2.0"), reason="TE 2.2.0 is required") + @pytest.mark.parametrize("tp_size", [2]) + @pytest.mark.parametrize("dp_overlap", [(True, True)]) + @pytest.mark.skipif(not cuda_graph_supported, reason=reason_for_no_cuda_graph) + def test_tensorwise_scaling_with_cuda_graph(self, tp_size, dp_overlap): + kwargs = {"overlap_param_gather": dp_overlap[0], "overlap_grad_reduce": dp_overlap[1]} + self.run_test_with_cuda_graph(tp_size, "tensorwise", **kwargs) + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) @pytest.mark.skipif(not is_te_min_version("2.2.0"), reason="TE 2.2.0 is required") @pytest.mark.parametrize("tp_size", [2]) @@ -282,6 +337,18 @@ def test_blockwise_scaling(self, tp_size, dp_overlap): kwargs = {"overlap_param_gather": dp_overlap[0], "overlap_grad_reduce": dp_overlap[1]} self.run_test(tp_size=tp_size, recipe="blockwise") + @pytest.mark.skipif( + get_device_arch_version() != 9, reason="blockwise is only supported on Hopper architecture" + ) + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + @pytest.mark.skipif(not is_te_min_version("2.4.0.dev0"), reason="TE 2.4.0.dev0 is required") + @pytest.mark.parametrize("tp_size", [2]) + @pytest.mark.parametrize("dp_overlap", [(True, True)]) + @pytest.mark.skipif(not cuda_graph_supported, reason=reason_for_no_cuda_graph) + def test_blockwise_scaling_with_cuda_graph(self, tp_size, dp_overlap): + kwargs = {"overlap_param_gather": dp_overlap[0], "overlap_grad_reduce": dp_overlap[1]} + self.run_test_with_cuda_graph(tp_size, "blockwise", **kwargs) + @pytest.mark.skipif( get_device_arch_version() < 10, reason="MXFP8 is supported since Blackwell architecture" ) @@ -296,6 +363,21 @@ def test_mxfp8(self, tp_size, dp_overlap): kwargs = {"overlap_param_gather": dp_overlap[0], "overlap_grad_reduce": dp_overlap[1]} self.run_test(tp_size=tp_size, recipe="mxfp8", **kwargs) + @pytest.mark.skipif( + get_device_arch_version() < 10, reason="MXFP8 is supported since Blackwell architecture" + ) + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + @pytest.mark.skipif(not is_te_min_version("2.3.0.dev0"), reason="TE 2.3.0.dev0 is required") + @pytest.mark.parametrize("tp_size", [2]) + @pytest.mark.parametrize("dp_overlap", [(False, False), (False, True), (True, True)]) + @pytest.mark.skipif(not cuda_graph_supported, reason=reason_for_no_cuda_graph) + def test_mxfp8_with_cuda_graph(self, tp_size, dp_overlap): + """ + dp_overlap: (overlap_param_gather, overlap_grad_reduce) + """ + kwargs = {"overlap_param_gather": dp_overlap[0], "overlap_grad_reduce": dp_overlap[1]} + self.run_test_with_cuda_graph(tp_size=tp_size, recipe="mxfp8", **kwargs) + @pytest.mark.skipif( get_device_arch_version() != 9, reason="blockwise is only supported on Hopper architecture" ) From 3b96f30e1ec281b7e1d39047038d9ad4349261ad Mon Sep 17 00:00:00 2001 From: kunlunl Date: Fri, 24 Oct 2025 13:45:25 +0800 Subject: [PATCH 3/9] Fix doc string & cuda graph scope --- megatron/core/fp8_utils.py | 5 +++++ tests/unit_tests/test_fp8_param.py | 1 - 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/megatron/core/fp8_utils.py b/megatron/core/fp8_utils.py index 18f0187953a..761616bc5ac 100644 --- a/megatron/core/fp8_utils.py +++ b/megatron/core/fp8_utils.py @@ -407,6 +407,11 @@ def correct_amax_history_if_needed(model: List[torch.nn.Module]): def post_all_gather_processing(model_params): + """ + Post-processing after all-gather for weights in distributed optimizer. + - tensorwise: may need to create a transposed view to match backend GEMM. + - blockwise: create column-wise storage. + """ try: from transformer_engine.pytorch.tensor.utils import post_all_gather_processing diff --git a/tests/unit_tests/test_fp8_param.py b/tests/unit_tests/test_fp8_param.py index 15cc17b89a9..217ec0ebc6a 100644 --- a/tests/unit_tests/test_fp8_param.py +++ b/tests/unit_tests/test_fp8_param.py @@ -121,7 +121,6 @@ def create_test_args( if use_cuda_graph: args.external_cuda_graph = True - args.cuda_graph_scope = "attn" args.cuda_graph_warmup_steps = 0 for key, value in kwargs.items(): From 040e827ef957c950fc932ffdb434716351d19689 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Mon, 3 Nov 2025 15:33:37 +0800 Subject: [PATCH 4/9] Update CUDA graph arg --- tests/unit_tests/test_fp8_param.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/test_fp8_param.py b/tests/unit_tests/test_fp8_param.py index 217ec0ebc6a..42386a87ce2 100644 --- a/tests/unit_tests/test_fp8_param.py +++ b/tests/unit_tests/test_fp8_param.py @@ -120,7 +120,7 @@ def create_test_args( args.reuse_grad_buf_for_mxfp8_param_ag = True if use_cuda_graph: - args.external_cuda_graph = True + args.cuda_graph_impl = "transformer_engine" args.cuda_graph_warmup_steps = 0 for key, value in kwargs.items(): From d450bbee52dcdcb07df365bb0c4c7ac360724933 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Tue, 4 Nov 2025 18:17:10 +0800 Subject: [PATCH 5/9] Really eanble cuda graph in UT --- megatron/core/fp8_utils.py | 27 ++++++++--- .../core/transformer/transformer_config.py | 2 +- tests/unit_tests/test_fp8_param.py | 48 +++++++++++++++++++ 3 files changed, 70 insertions(+), 7 deletions(-) diff --git a/megatron/core/fp8_utils.py b/megatron/core/fp8_utils.py index 761616bc5ac..d7f4f0c858b 100644 --- a/megatron/core/fp8_utils.py +++ b/megatron/core/fp8_utils.py @@ -78,6 +78,13 @@ Fp8Padding = None Fp8Unpadding = None +try: + from transformer_engine.pytorch.tensor.utils import ( + post_all_gather_processing as te_post_all_gather_processing, + ) +except ImportError: + te_post_all_gather_processing = None + def is_float8tensor(tensor: torch.Tensor) -> bool: """Check if a tensor is a Transformer Engine Float8Tensor. @@ -172,7 +179,15 @@ def _quantize_param_shard_impl( raise NotImplementedError( f"FSDP with --fp8-param-gather is not supported in TE v{get_te_version()}" ) - cast_master_weights_to_fp8(*args) + + # For newer TE versions (i.e., have post_all_gather_processing function), we keep the + # columnwise data and manually call post_all_gather_processing after all-gather, this + # makes fp8 params compatible with CUDA graph. + kwargs = {} + if te_post_all_gather_processing is not None: + kwargs["keep_columnwise"] = True + + cast_master_weights_to_fp8(*args, **kwargs) def _correct_amax_history_if_needed_impl(model: List[torch.nn.Module]) -> None: pass @@ -412,11 +427,11 @@ def post_all_gather_processing(model_params): - tensorwise: may need to create a transposed view to match backend GEMM. - blockwise: create column-wise storage. """ - try: - from transformer_engine.pytorch.tensor.utils import post_all_gather_processing - - post_all_gather_processing(model_params) - except ImportError: + if te_post_all_gather_processing is not None: + te_post_all_gather_processing(model_params) + else: + # If the TE version is old and does not have post_all_gather_processing function, this is + # a no-op, and the transpose/columnwise data will be created in the next forward pass. pass diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 63d0d2efd27..aab137b6430 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -693,7 +693,7 @@ class TransformerConfig(ModelParallelConfig): cuda_graph_scope: Optional[List[str]] = None """Determines the CUDA graphs capturing scope. When cuda_graph_impl is set to "transformer_engine", valid values are "attn", "mlp", "moe", - "moe_router", "moe_preprocess", "mamba". None means ["attn", "mlp"]. + "moe_router", "moe_preprocess", "mamba". None means the full layer. When cuda_graph_impl is set to "local", "full_iteration" can be specified as cuda_graph_scope to enable whole iteration CUDA graph. All other values enable layerwise CUDA graph.""" diff --git a/tests/unit_tests/test_fp8_param.py b/tests/unit_tests/test_fp8_param.py index 42386a87ce2..eee6a3bbeb4 100644 --- a/tests/unit_tests/test_fp8_param.py +++ b/tests/unit_tests/test_fp8_param.py @@ -8,6 +8,7 @@ import torch from transformer_engine.pytorch.fp8 import check_fp8_support +from megatron.core.distributed import DistributedDataParallel as DDP from megatron.core.enums import ModelType from megatron.core.fp8_utils import is_float8tensor from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec @@ -39,6 +40,27 @@ reason_for_no_cuda_graph = "Need newer TransformerEngine" +def enable_forward_pre_hook(model_chunks): + for model_chunk in model_chunks: + assert isinstance(model_chunk, DDP) + model_chunk.enable_forward_pre_hook() + + +def disable_forward_pre_hook(model_chunks, param_sync=True): + for model_chunk in model_chunks: + assert isinstance(model_chunk, DDP) + model_chunk.disable_forward_pre_hook(param_sync=param_sync) + + +def should_disable_forward_pre_hook(args): + """Block forward pre-hook for certain configurations.""" + return ( + not args.use_megatron_fsdp + and args.use_distributed_optimizer + and args.overlap_param_gather + ) + + class TestFP8Param: def setup_method(self, method): @@ -122,6 +144,7 @@ def create_test_args( if use_cuda_graph: args.cuda_graph_impl = "transformer_engine" args.cuda_graph_warmup_steps = 0 + args.use_te_rng_tracker = args.cuda_graph_impl != "none" for key, value in kwargs.items(): assert hasattr(args, key) @@ -186,6 +209,20 @@ def _run_test_helper( ) assert len(gpt_model) == 1 # Assume only one model in the model provider. + cuda_graph_helper = None + # Hard coded to use cuda_graph_impl="transformer_engine" + cuda_graph_impl = "transformer_engine" + if use_cuda_graph and cuda_graph_impl == "transformer_engine": + from megatron.core.transformer.cuda_graphs import TECudaGraphHelper + + cuda_graph_helper = TECudaGraphHelper( + model=gpt_model, + config=gpt_model[0].config, + seq_length=self.seq_length, + micro_batch_size=self.micro_batch_size, + optimizers=[optimizer], + ) + num_fp8_params = 0 for _, param in gpt_model[0].named_parameters(): if not inference: @@ -210,6 +247,17 @@ def _run_test_helper( gpt_model[0].zero_grad_buffer() optimizer.zero_grad() + # Capture CUDA graphs after warmup if helper is provided. + # Hard coded cuda_graph_warmup_steps = 0. + cuda_graph_warmup_steps = 0 + if cuda_graph_helper is not None and i == cuda_graph_warmup_steps: + if should_disable_forward_pre_hook(args): + disable_forward_pre_hook(gpt_model, param_sync=False) + cuda_graph_helper.create_cudagraphs() + if should_disable_forward_pre_hook(args): + enable_forward_pre_hook(gpt_model) + cuda_graph_helper.cuda_graph_set_manual_hooks() + # For the mxfp8_param with reuse_grad_buf_for_mxfp8_param_ag and dp_ag_overlap, # we need to call the _copy_main_params_to_param_buffer() after the grad buffer # is zeroed by zero_grad_buffer() because param and grad buffer are shared. From ffca6f082599fa70de9da0a94bcd3c305da1d94d Mon Sep 17 00:00:00 2001 From: kunlunl Date: Wed, 5 Nov 2025 15:28:28 +0800 Subject: [PATCH 6/9] Rename keep_columnwise to manual_post_all_gather_processing Signed-off-by: kunlunl --- megatron/core/fp8_utils.py | 2 +- tests/unit_tests/test_fp8_param.py | 2 ++ tests/unit_tests/transformer/test_cuda_graphs.py | 2 ++ 3 files changed, 5 insertions(+), 1 deletion(-) diff --git a/megatron/core/fp8_utils.py b/megatron/core/fp8_utils.py index d7f4f0c858b..7c3591ae5f7 100644 --- a/megatron/core/fp8_utils.py +++ b/megatron/core/fp8_utils.py @@ -185,7 +185,7 @@ def _quantize_param_shard_impl( # makes fp8 params compatible with CUDA graph. kwargs = {} if te_post_all_gather_processing is not None: - kwargs["keep_columnwise"] = True + kwargs["manual_post_all_gather_processing"] = True cast_master_weights_to_fp8(*args, **kwargs) diff --git a/tests/unit_tests/test_fp8_param.py b/tests/unit_tests/test_fp8_param.py index eee6a3bbeb4..f64d4f648ef 100644 --- a/tests/unit_tests/test_fp8_param.py +++ b/tests/unit_tests/test_fp8_param.py @@ -3,6 +3,7 @@ import contextlib import os import sys +import gc import pytest import torch @@ -72,6 +73,7 @@ def teardown_method(self, method): Utils.destroy_model_parallel() destroy_global_vars() destroy_num_microbatches_calculator() + gc.collect() def model_provider( self, diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index 1302369266a..0fc6598e64b 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/tests/unit_tests/transformer/test_cuda_graphs.py @@ -5,6 +5,7 @@ import sys import time import types +import gc import pytest import torch @@ -763,6 +764,7 @@ def teardown_method(self, method): Utils.destroy_model_parallel() destroy_global_vars() destroy_num_microbatches_calculator() + gc.collect() def model_provider( self, From becb3fbe7aee2e251d6851064a6a50e90cf0aa91 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Fri, 7 Nov 2025 14:21:00 +0800 Subject: [PATCH 7/9] Fix lint error Signed-off-by: kunlunl --- tests/unit_tests/test_fp8_param.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/unit_tests/test_fp8_param.py b/tests/unit_tests/test_fp8_param.py index f64d4f648ef..d7ae588131c 100644 --- a/tests/unit_tests/test_fp8_param.py +++ b/tests/unit_tests/test_fp8_param.py @@ -56,9 +56,7 @@ def disable_forward_pre_hook(model_chunks, param_sync=True): def should_disable_forward_pre_hook(args): """Block forward pre-hook for certain configurations.""" return ( - not args.use_megatron_fsdp - and args.use_distributed_optimizer - and args.overlap_param_gather + not args.use_megatron_fsdp and args.use_distributed_optimizer and args.overlap_param_gather ) From 0d79b3ac4d1ff7d36973c62e030588122d36ec1a Mon Sep 17 00:00:00 2001 From: kunlunl Date: Fri, 7 Nov 2025 16:09:12 +0800 Subject: [PATCH 8/9] Fix lint error Signed-off-by: kunlunl --- tests/unit_tests/test_fp8_param.py | 2 +- tests/unit_tests/transformer/test_cuda_graphs.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/test_fp8_param.py b/tests/unit_tests/test_fp8_param.py index d7ae588131c..b318d0c0afa 100644 --- a/tests/unit_tests/test_fp8_param.py +++ b/tests/unit_tests/test_fp8_param.py @@ -1,9 +1,9 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. import contextlib +import gc import os import sys -import gc import pytest import torch diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index 0fc6598e64b..fb3567074f1 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/tests/unit_tests/transformer/test_cuda_graphs.py @@ -1,11 +1,11 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +import gc import os import random import sys import time import types -import gc import pytest import torch From 63fd2755ad9730480241f9a54eebbde665a1492e Mon Sep 17 00:00:00 2001 From: kunlunl Date: Wed, 12 Nov 2025 15:06:12 +0800 Subject: [PATCH 9/9] Remove invalid use_te_rng_tracker arg Signed-off-by: kunlunl --- tests/unit_tests/test_fp8_param.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/test_fp8_param.py b/tests/unit_tests/test_fp8_param.py index b318d0c0afa..0b8d41769ec 100644 --- a/tests/unit_tests/test_fp8_param.py +++ b/tests/unit_tests/test_fp8_param.py @@ -144,7 +144,6 @@ def create_test_args( if use_cuda_graph: args.cuda_graph_impl = "transformer_engine" args.cuda_graph_warmup_steps = 0 - args.use_te_rng_tracker = args.cuda_graph_impl != "none" for key, value in kwargs.items(): assert hasattr(args, key) @@ -194,6 +193,7 @@ def _run_test_helper( set_args(args) torch.manual_seed(_SEED) Utils.initialize_model_parallel(tensor_model_parallel_size=tp_size) + input_ids, labels, position_ids, attention_mask, loss_mask = self.get_batch( self.seq_length, self.micro_batch_size )