diff --git a/src/megatron/bridge/models/model_provider.py b/src/megatron/bridge/models/model_provider.py index 0330e61f2f..ccf304d13d 100644 --- a/src/megatron/bridge/models/model_provider.py +++ b/src/megatron/bridge/models/model_provider.py @@ -677,17 +677,18 @@ def _ddp_wrap( else: DP = DistributedDataParallel - model = [ - DP( - config=get_model_config(model_chunk), - ddp_config=ddp_config, - module=model_chunk, - # Turn off bucketing for model_chunk 2 onwards, since communication for these - # model chunks is overlapped with compute anyway. - disable_bucketing=(model_chunk_idx > 0) or overlap_param_gather_with_optimizer_step, - ) - for (model_chunk_idx, model_chunk) in enumerate(model) - ] + with torch.cuda.stream(torch.cuda.Stream()): + model = [ + DP( + config=get_model_config(model_chunk), + ddp_config=ddp_config, + module=model_chunk, + # Turn off bucketing for model_chunk 2 onwards, since communication for these + # model chunks is overlapped with compute anyway. + disable_bucketing=(model_chunk_idx > 0) or overlap_param_gather_with_optimizer_step, + ) + for (model_chunk_idx, model_chunk) in enumerate(model) + ] # Broadcast params from data parallel src rank to other data parallel ranks. if data_parallel_random_init: diff --git a/src/megatron/bridge/recipes/qwen/qwen3_moe.py b/src/megatron/bridge/recipes/qwen/qwen3_moe.py index a1d6691e47..926512f8c7 100644 --- a/src/megatron/bridge/recipes/qwen/qwen3_moe.py +++ b/src/megatron/bridge/recipes/qwen/qwen3_moe.py @@ -225,8 +225,6 @@ def _qwen3_moe_common( if precision_config is None: precision_config = bf16_mixed() - if isinstance(precision_config, MixedPrecisionConfig): - precision_config.grad_reduce_in_fp32 = False # MoE-specific pipeline split configurations if account_for_embedding_in_pipeline_split: