File size: 4,700 Bytes
ed12f96 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | diff --git a/src/megatron/bridge/models/conversion/param_mapping.py b/src/megatron/bridge/models/conversion/param_mapping.py
index dc7d0be..8156826 100644
--- a/src/megatron/bridge/models/conversion/param_mapping.py
+++ b/src/megatron/bridge/models/conversion/param_mapping.py
@@ -1088,15 +1088,19 @@ class AutoMapping(MegatronParamMapping[torch.Tensor]):
"ColumnParallelLinear",
"TEColumnParallelLinear",
"TELayerNormColumnParallelLinear",
+ "MindSpeedTELayerNormColumnParallelLinear",
"TEColumnParallelGroupedLinear",
+ "MindSpeedTEColumnParallelGroupedLinear",
"VocabParallelEmbedding",
"DotProductAttention", # for attention sink only
"TEDotProductAttention", # for attention sink only
+ "MindSpeedTEDotProductAttention",
},
"row": {
"RowParallelLinear",
"TERowParallelLinear",
"TERowParallelGroupedLinear",
+ "MindSpeedTERowParallelGroupedLinear",
},
"replicated": {
# Normalization layers
@@ -1164,7 +1168,7 @@ class AutoMapping(MegatronParamMapping[torch.Tensor]):
# Handle fused modules like TELayerNormColumnParallelLinear
# These modules have both column-parallel weights (weight, bias)
# and replicated layer norm weights (layer_norm_weight, layer_norm_bias)
- if module_type == "TELayerNormColumnParallelLinear":
+ if module_type == "TELayerNormColumnParallelLinear" or module_type == "MindSpeedTELayerNormColumnParallelLinear":
# Check the actual parameter name to determine the correct parallelism type
if self.megatron_param and (
self.megatron_param.endswith("layer_norm_weight") or self.megatron_param.endswith("layer_norm_bias")
@@ -1195,7 +1199,7 @@ class AutoMapping(MegatronParamMapping[torch.Tensor]):
return "replicated"
# Check parallel_mode for TELinear
- if module_type == "TELinear":
+ if module_type == "TELinear" or module_type == "MindSpeedTELinear":
if module.parallel_mode == "column":
return "column"
elif module.parallel_mode == "row":
diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_block.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_block.py
index 3f64d8d..9dadc4b 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_block.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_block.py
@@ -71,8 +71,9 @@ class Qwen3VLTransformerBlock(TransformerBlock):
context_mask,
rotary_pos_emb,
visual_pos_masks,
- deepstack_visual_embeds,
+ *deepstack_visual_embeds_args,
):
+ deepstack_visual_embeds = list(deepstack_visual_embeds_args) if deepstack_visual_embeds_args else None
for index in range(start, end):
layer = self._get_layer(index)
inner_fp8_context = (
@@ -103,6 +104,8 @@ class Qwen3VLTransformerBlock(TransformerBlock):
return hidden_states, context
return custom_forward
+
+ deepstack_visual_embeds_tuple = tuple(deepstack_visual_embeds) if deepstack_visual_embeds else ()
def checkpoint_handler(forward_func):
"""Determines whether to use the `te_checkpoint` or `tensor_parallel.checkpoint`"""
@@ -118,7 +121,7 @@ class Qwen3VLTransformerBlock(TransformerBlock):
context_mask,
rotary_pos_emb,
visual_pos_masks,
- deepstack_visual_embeds,
+ *deepstack_visual_embeds_tuple,
)
else:
return tensor_parallel.checkpoint(
@@ -130,7 +133,7 @@ class Qwen3VLTransformerBlock(TransformerBlock):
context_mask,
rotary_pos_emb,
visual_pos_masks,
- deepstack_visual_embeds,
+ *deepstack_visual_embeds_tuple,
)
if self.config.recompute_method == "uniform":
@@ -169,7 +172,7 @@ class Qwen3VLTransformerBlock(TransformerBlock):
context_mask,
rotary_pos_emb,
visual_pos_masks,
- deepstack_visual_embeds,
+ *deepstack_visual_embeds_tuple,
)
else:
raise ValueError("Invalid activation recompute method.")
|