@@ -86,7 +86,7 @@ class TorchCommonLoadStrategy(LoadCommonStrategy):
msc = MultiStorageClientFeature.import_package()
return msc.torch.load(load_path, map_location='cpu')
else:
- return torch.load(load_path, map_location='cpu')
+ return torch.load(load_path, map_location='cpu', weights_only=False)
except FileNotFoundError as e:
err_msg = f'Common file {load_path} does not exist'
if MultiStorageClientFeature.is_enabled():
@@ -427,6 +427,15 @@ def _restore_dict_types(x: Union[dict, list, Any], keys_template: Union[dict, li
_restore_dict_types(x_val, templ_val)
+@dataclass
+class MCoreMetadata(Metadata):
+ """Metadata with mcore specific data."""
+
+ # holds data related to flattened_range
+ # TODO: remove when flattened_range is properly removed
+ mcore_data: Optional[Dict[str, Dict[str, Any]]] = None # Mcore related data about each tensor
+
+
@dataclass(frozen=True)
class MCoreSavePlan(SavePlan):
"""SavePlan with MCore specific data."""
@@ -499,9 +508,10 @@ class MCoreSavePlanner(DefaultSavePlanner):
def create_global_plan(self, all_plans: List[MCoreSavePlan]) -> Tuple[List[SavePlan], Metadata]:
"""Merges MCore data for all plans."""
global_plan, metadata = super().create_global_plan(all_plans)
- metadata.mcore_data = dict(
+ mcore_data = dict(
ChainMap(*(plan.mcore_data for plan in all_plans)) # type: ignore[arg-type]
)
+ metadata = MCoreMetadata(mcore_data=mcore_data, **vars(metadata))
return global_plan, metadata
def create_decentralized_global_plan(self, local_plan: SavePlan) -> SavePlan:
@@ -556,10 +566,12 @@ class MCoreLoadPlanner(DefaultLoadPlanner):
def _validate_global_shapes(self, metadata, sharded_tensors):
for sh_ten in sharded_tensors:
if sh_ten.key not in metadata.state_dict_metadata:
- raise KeyError(
- f"{sh_ten.key} from model not in state dict:"
- f" {sorted(metadata.state_dict_metadata.keys())}"
- )
+ # raise KeyError(
+ # f"{sh_ten.key} from model not in state dict:"
+ # f" {sorted(metadata.state_dict_metadata.keys())}"
+ # )
+ print(f"{sh_ten.key} from model not in state dict, will skip")
+ continue
loaded_shape = metadata.state_dict_metadata[sh_ten.key].size
expected_shape = self._expected_shape(sh_ten)
if loaded_shape != expected_shape:
@@ -589,7 +601,7 @@ class MCoreLoadPlanner(DefaultLoadPlanner):
tensor_metadata = self.metadata.state_dict_metadata
metadata_with_sizes = [
(tensor_metadata[key], tensor_metadata[key].size, sharded_tensor)
- for key, sharded_tensor in self.allow_shape_mismatch_sharded_tensors.items()
+ for key, sharded_tensor in self.allow_shape_mismatch_sharded_tensors.items() if key in tensor_metadata
]
try:
# Temporarily set sizes to expected shapes
@@ -918,6 +930,7 @@ class TorchDistLoadShardedStrategy(LoadShardedStrategy):
planner=MCoreLoadPlanner(
shapes_validation_sharded_tensors=flexible_shape_sharded_tensors,
allow_shape_mismatch_sharded_tensors=allow_shape_mismatch_sharded_tensors,
+ allow_partial_load=True,
),
)
@@ -11,3 +11,15 @@ from .finalize_model_grads import finalize_model_grads
from .fsdp.mcore_fsdp_adapter import FullyShardedDataParallel
from .torch_fully_sharded_data_parallel import TorchFullyShardedDataParallel
from .torch_fully_sharded_data_parallel_config import TorchFullyShardedDataParallelConfig
+
+# Backward compatibility patch for FSDP module reorganization
+import sys
+import importlib.util
+
+spec = importlib.util.find_spec('megatron.core.distributed.fsdp.src.megatron_fsdp')
+if spec:
+ custom_fsdp = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(custom_fsdp)
+ sys.modules['megatron.core.distributed.custom_fsdp'] = custom_fsdp
+ if hasattr(custom_fsdp, 'MegatronFSDP'):
+ custom_fsdp.FullyShardedDataParallel = custom_fsdp.MegatronFSDP
@@ -366,6 +366,7 @@ class TELinear(te.pytorch.Linear):
)
for param in self.parameters():
+ setattr(param, "parallel_mode", parallel_mode)
if is_expert:
# Reduce the gradient on the expert_data_parallel group for expert linear layers
setattr(param, "allreduce", not self.expert_parallel)
@@ -79,6 +79,8 @@ def get_gpt_layer_with_transformer_engine_spec(
qk_l2_norm: Optional[bool] = False,
use_te_op_fuser: Optional[bool] = False,
use_kitchen: bool = False,
+ post_self_attn_layernorm: bool = False,
+ post_mlp_layernorm: bool = False,
) -> ModuleSpec:
"""Use this spec to use lower-level Transformer Engine modules (required for fp8 training).
@@ -178,9 +180,11 @@ def get_gpt_layer_with_transformer_engine_spec(
),
),
self_attn_bda=get_bias_dropout_add,
+ post_self_attn_layernorm=TENorm if post_self_attn_layernorm else IdentityOp,
pre_mlp_layernorm=backend.layer_norm() if num_experts else IdentityOp,
mlp=mlp,
mlp_bda=get_bias_dropout_add,
+ post_mlp_layernorm=TENorm if post_mlp_layernorm else IdentityOp,
sharded_state_dict_keys_map={
"mlp.0.weight": "mlp.linear_fc1.layer_norm_weight",
"mlp.0.bias": "mlp.linear_fc1.layer_norm_bias",
@@ -355,6 +355,7 @@ class GPTModel(LanguageModule):
*,
inference_params: Optional[BaseInferenceContext] = None,
loss_mask: Optional[Tensor] = None,
+ mtp_kwargs: Optional[dict] = {},
) -> Tensor:
"""Forward function of the GPT Model This function passes the input tensors
through the embedding layer, and then the decoeder and finally into the post
@@ -410,6 +411,7 @@ class GPTModel(LanguageModule):
runtime_gather_output=runtime_gather_output,
extra_block_kwargs=extra_block_kwargs,
inference_context=inference_context,
+ mtp_kwargs=mtp_kwargs,
)
def _postprocess(
@@ -431,6 +433,7 @@ class GPTModel(LanguageModule):
runtime_gather_output=None,
extra_block_kwargs=None,
inference_context=None,
+ mtp_kwargs={},
):
"""Postprocesses decoder hidden states to generate logits or compute loss.
@@ -446,7 +449,7 @@ class GPTModel(LanguageModule):
if self.share_embeddings_and_output_weights:
output_weight = self.shared_embedding_or_output_weight()
- if mtp_in_postprocess:
+ if mtp_in_postprocess and mtp_kwargs.get('mtp_labels', None) is not None:
hidden_states = self.mtp(
input_ids=input_ids,
position_ids=position_ids,
@@ -465,25 +468,37 @@ class GPTModel(LanguageModule):
if not self.post_process:
return hidden_states
- if self.mtp_process:
- mtp_labels = labels.clone()
+ if self.mtp_process and mtp_kwargs.get('mtp_labels', None) is not None:
+ mtp_labels = mtp_kwargs['mtp_labels'].clone()
+ mtp_labels, _ = roll_tensor(mtp_labels, shifts=-1, dims=-1, cp_group=self.cp_group, packed_seq_params=packed_seq_params)
+
hidden_states_list = torch.chunk(hidden_states, 1 + self.config.mtp_num_layers, dim=0)
hidden_states = hidden_states_list[0]
if loss_mask is None:
# if loss_mask is not provided, use all ones as loss_mask
loss_mask = torch.ones_like(mtp_labels)
+ else:
+ # Otherwise, roll the loss_mask to keep up with the mtp_labels
+ loss_mask, _ = roll_tensor(loss_mask, shifts=-1, dims=-1, cp_group=self.cp_group, packed_seq_params=packed_seq_params)
for mtp_layer_number in range(self.config.mtp_num_layers):
# output
- mtp_logits, _ = self.output_layer(
- hidden_states_list[mtp_layer_number + 1],
- weight=output_weight,
- runtime_gather_output=runtime_gather_output,
+ output_layer_params = {k: v.detach() for k, v in self.output_layer.named_parameters()}
+ output_layer_buffers = dict(self.output_layer.named_buffers())
+ mtp_logits, _ = torch.func.functional_call(
+ self.output_layer,
+ {**output_layer_params, **output_layer_buffers},
+ (hidden_states_list[mtp_layer_number + 1],),
+ {
+ "weight": output_weight.detach() if output_weight else None,
+ "runtime_gather_output": runtime_gather_output,
+ },
)
# Calc loss for the current Multi-Token Prediction (MTP) layers.
- mtp_labels, _ = roll_tensor(mtp_labels, shifts=-1, dims=-1, cp_group=self.cp_group)
- loss_mask, num_tokens = roll_tensor(
- loss_mask, shifts=-1, dims=-1, cp_group=self.cp_group
+ mtp_labels, _ = roll_tensor(mtp_labels, shifts=-1, dims=-1, cp_group=self.cp_group, packed_seq_params=packed_seq_params)
+ new_loss_mask, num_tokens = roll_tensor(
+ loss_mask, shifts=-1, dims=-1, cp_group=self.cp_group, packed_seq_params=packed_seq_params
)
+ loss_mask = new_loss_mask * loss_mask
mtp_loss = self.compute_language_model_loss(mtp_labels, mtp_logits)
mtp_loss = loss_mask * mtp_loss
if self.training:
@@ -9,6 +9,7 @@ from typing import Callable, List, Optional
import numpy as np
import torch
+import torch.distributed as dist
from .utils import GlobalMemoryBuffer, is_torch_min_version
@@ -26,22 +26,22 @@ def _batched_p2p_ops(
ops = []
if tensor_send_prev is not None:
send_prev_op = torch.distributed.P2POp(
- torch.distributed.isend, tensor_send_prev, prev_pipeline_rank, group
+ torch.distributed.isend, tensor_send_prev, prev_pipeline_rank,
)
ops.append(send_prev_op)
if tensor_recv_prev is not None:
recv_prev_op = torch.distributed.P2POp(
- torch.distributed.irecv, tensor_recv_prev, prev_pipeline_rank, group
+ torch.distributed.irecv, tensor_recv_prev, prev_pipeline_rank,
)
ops.append(recv_prev_op)
if tensor_send_next is not None:
send_next_op = torch.distributed.P2POp(
- torch.distributed.isend, tensor_send_next, next_pipeline_rank, group
+ torch.distributed.isend, tensor_send_next, next_pipeline_rank,
)
ops.append(send_next_op)
if tensor_recv_next is not None:
recv_next_op = torch.distributed.P2POp(
- torch.distributed.irecv, tensor_recv_next, next_pipeline_rank, group
+ torch.distributed.irecv, tensor_recv_next, next_pipeline_rank,
)
ops.append(recv_next_op)
if len(ops) > 0:
@@ -670,7 +670,10 @@ class Attention(MegatronModule, ABC):
# Get the query, key and value tensors based on the type of attention -
# self or cross attn.
nvtx_range_push(suffix="qkv")
- query, key, value = self.get_query_key_value_tensors(hidden_states, key_value_states)
+ if self.config.use_gated_attention:
+ query, gate, key, value = self.get_query_gate_key_value_tensors(hidden_states, key_value_states)
+ else:
+ query, key, value = self.get_query_key_value_tensors(hidden_states, key_value_states)
nvtx_range_pop(suffix="qkv")
#
@@ -842,6 +845,11 @@ class Attention(MegatronModule, ABC):
# Output. [sq, b, h]
#
+ if self.config.use_gated_attention:
+ nvtx_range_push(suffix="sigmoid_gate")
+ core_attn_out = core_attn_out * torch.sigmoid(gate)
+ nvtx_range_pop(suffix="sigmoid_gate")
+
nvtx_range_push(suffix="linear_proj")
output, bias = self.linear_proj(core_attn_out)
nvtx_range_pop(suffix="linear_proj")
@@ -879,19 +887,34 @@ class SelfAttention(Attention):
model_comm_pgs=model_comm_pgs,
)
- self.linear_qkv = build_module(
- submodules.linear_qkv,
- self.config.hidden_size,
- self.query_projection_size + 2 * self.kv_projection_size,
- config=self.config,
- init_method=self.config.init_method,
- gather_output=False,
- bias=self.config.add_bias_linear or self.config.add_qkv_bias,
- skip_bias_add=False,
- is_expert=False,
- tp_comm_buffer_name='qkv',
- tp_group=self.model_comm_pgs.tp,
- )
+ if self.config.use_gated_attention:
+ self.linear_qgkv = build_module(
+ submodules.linear_qkv,
+ self.config.hidden_size,
+ 2 * (self.query_projection_size + self.kv_projection_size),
+ config=self.config,
+ init_method=self.config.init_method,
+ gather_output=False,
+ bias=self.config.add_bias_linear or self.config.add_qkv_bias,
+ skip_bias_add=False,
+ is_expert=False,
+ tp_comm_buffer_name='qkv',
+ tp_group=self.model_comm_pgs.tp,
+ )
+ else:
+ self.linear_qkv = build_module(
+ submodules.linear_qkv,
+ self.config.hidden_size,
+ self.query_projection_size + 2 * self.kv_projection_size,
+ config=self.config,
+ init_method=self.config.init_method,
+ gather_output=False,
+ bias=self.config.add_bias_linear or self.config.add_qkv_bias,
+ skip_bias_add=False,
+ is_expert=False,
+ tp_comm_buffer_name='qkv',
+ tp_group=self.model_comm_pgs.tp,
+ )
if submodules.q_layernorm is not None:
self.q_layernorm = build_module(
@@ -1036,6 +1059,65 @@ class SelfAttention(Attention):
return query, key, value
+ # adapt from https://github.com/alibaba/Pai-Megatron-Patch/blob/8e6cbb0556ba09933ab4a4edb23c0af1d19d9960/megatron_patch/model/qwen3_next/gated_attention.py#L192
+ def get_query_gate_key_value_tensors(self, hidden_states, key_value_states=None):
+ """
+ Derives `query`, `key` and `value` tensors from `hidden_states`.
+ """
+ # Attention heads [sq, b, h] --> [sq, b, ng * 2 * (np/ng + 1) * hn)]
+ mixed_qgkv, _ = self.linear_qgkv(hidden_states)
+
+ # [sq, b, hp] --> [sq, b, ng, 2 * (np/ng + 1) * hn]
+ new_tensor_shape = mixed_qgkv.size()[:-1] + (
+ self.num_query_groups_per_partition,
+ (
+ 2 * (self.num_attention_heads_per_partition // self.num_query_groups_per_partition + 1)
+ * self.hidden_size_per_attention_head
+ ),
+ )
+ mixed_qgkv = mixed_qgkv.view(*new_tensor_shape)
+
+ split_arg_list = [
+ (
+ self.num_attention_heads_per_partition
+ // self.num_query_groups_per_partition
+ * self.hidden_size_per_attention_head
+ ),
+ (
+ self.num_attention_heads_per_partition
+ // self.num_query_groups_per_partition
+ * self.hidden_size_per_attention_head
+ ),
+ self.hidden_size_per_attention_head,
+ self.hidden_size_per_attention_head,
+ ]
+
+ if SplitAlongDim is not None:
+
+ # [sq, b, ng, (np/ng + 2) * hn]
+ # --> [sq, b, ng, np/ng * hn], [sq, b, ng, hn], [sq, b, ng, hn]
+ (query, gate, key, value) = SplitAlongDim(mixed_qgkv, 3, split_arg_list)
+ else:
+
+ # [sq, b, ng, (np/ng + 2) * hn]
+ # --> [sq, b, ng, np/ng * hn], [sq, b, ng, hn], [sq, b, ng, hn]
+ (query, gate, key, value) = torch.split(mixed_qgkv, split_arg_list, dim=3)
+
+ # [sq, b, ng, np/ng * hn] -> [sq, b, np, hn]
+ query = query.reshape(query.size(0), query.size(1), -1, self.hidden_size_per_attention_head)
+ gate = gate.reshape(query.size(0), query.size(1), -1)
+
+ if self.q_layernorm is not None:
+ query = self.q_layernorm(query)
+
+ if self.k_layernorm is not None:
+ key = self.k_layernorm(key)
+
+ if self.config.test_mode:
+ self.run_realtime_tests()
+
+ return query, gate, key, value
+
def backward_dw(self) -> NoReturn:
"""Execute weight update operations"""
self._backward_qkv_proj()
@@ -566,6 +566,9 @@ def topk_routing_with_score_function(
else:
return torch.topk(scores, k=topk, dim=1)
+ from slime.utils.routing_replay import get_routing_replay_compute_topk
+ compute_topk = get_routing_replay_compute_topk(compute_topk)
+
if score_function == "softmax":
if use_pre_softmax:
scores = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(logits)
@@ -156,6 +156,9 @@ class TopKRouter(Router):
self.local_tokens_per_expert = None
self.expert_bias = None
+ from slime.utils.routing_replay import register_routing_replay
+ register_routing_replay(self)
+
def _maintain_float32_expert_bias(self):
"""
Maintain the expert bias in float32.
@@ -6,6 +6,7 @@ from typing import Callable, List, Optional, Union
import torch
from torch import Tensor
+import warnings
from megatron.core import InferenceParams, mpu, parallel_state, tensor_parallel
from megatron.core.dist_checkpointing.mapping import ShardedStateDict
@@ -105,17 +106,21 @@ def tie_output_layer_state_dict(
)
-def roll_tensor(tensor, shifts=-1, dims=-1, cp_group=None):
- """Roll the tensor input along the sequence dimension with Context Parallelism (CP) support.
- This function extends the original roll_tensor to support Context Parallelism, which allows
- MTP to work with CP > 1. When CP is enabled, the sequence dimension is split across CP ranks,
- and tensor rolling requires communication between adjacent CP ranks to properly handle the
- boundary conditions.
+def roll_tensor(tensor, shifts=-1, dims=-1, cp_group=None, packed_seq_params=None):
+ """Roll the tensor input along the sequence dimension with Context Parallelism (CP) and Packed Sequence support.
+
+ This function extends the original roll_tensor to support Context Parallelism and Packed Sequences.
+ When CP is enabled, the sequence dimension is split across CP ranks, and tensor rolling requires
+ communication between adjacent CP ranks to properly handle the boundary conditions.
+ When packed sequences are used, rolling is performed within each individual sequence boundary
+ to prevent mixing tokens between different packed sequences.
For CP=1 (default behavior): Uses standard torch.roll with zero padding
For CP>1: Splits tensor into chunks, performs rolling within each chunk, then exchanges
boundary elements between adjacent CP ranks to maintain sequence continuity.
+ For packed sequences: Rolls tensors within sequence boundaries defined by cu_seqlens.
+
Args:
tensor (Tensor): The input tensor to roll.
@@ -123,9 +128,15 @@ def roll_tensor(tensor, shifts=-1, dims=-1, cp_group=None):
dims (int): The dimension to roll (typically -1 for sequence dimension).
cp_group (ProcessGroup): The context parallelism process group. If None or size=1,
falls back to standard rolling behavior.
+ packed_seq_params (PackedSeqParams): Parameters for packed sequence processing.
+ If provided, rolling respects sequence boundaries.
Returns:
tuple: (rolled_tensor, sum_of_rolled_tensor)
"""
+
+ if packed_seq_params is not None:
+ return _roll_tensor_packed_seq(tensor, shifts, dims, packed_seq_params, cp_group)
+
# Standard rolling behavior when CP is not enabled (cp_group is None or size=1)
if cp_group is None or cp_group.size() == 1:
rolled_tensor = torch.roll(tensor, shifts=shifts, dims=dims)
@@ -193,6 +204,103 @@ def roll_tensor(tensor, shifts=-1, dims=-1, cp_group=None):
return rolled_tensor, rolled_tensor.sum()
+def _roll_tensor_packed_seq(tensor, shifts, dims, packed_seq_params, cp_group=None):
+ """Roll tensor with packed sequence support.
+
+ This function handles rolling for packed sequences by respecting sequence boundaries
+ defined in packed_seq_params.cu_seqlens. Rolling is performed within each individual
+ sequence to prevent mixing tokens between different packed sequences. When Context
+ Parallelism (CP) is enabled, each CP rank still receives the full `cu_seqlens` metadata
+ so we slice out the portion of every packed sequence that lives on the current rank and
+ reuse the standard CP boundary exchange to populate the rolling window.
+
+ Args:
+ tensor (Tensor): The input tensor to roll.
+ shifts (int): The shift of the tensor (typically -1 for MTP).
+ dims (int): The dimension to roll (typically -1 for sequence dimension).
+ packed_seq_params (PackedSeqParams): Parameters for packed sequence processing.
+ cp_group (ProcessGroup): The context parallelism process group.
+
+ Returns:
+ tuple: (rolled_tensor, sum_of_rolled_tensor)
+ """
+
+ # Notice: This is a naive implementation to test the correctness, a better solution will only sync the boundary tokens once.
+ assert dims == -1 or dims == tensor.dim() - 1, "Packed sequence roll only supports the last dimension."
+ assert shifts == -1, "Packed sequence roll only supports a single-token left shift."
+ cu_seqlens = packed_seq_params.cu_seqlens_q
+ assert cu_seqlens is not None, "Packed sequence parameters must provide cu_seqlens_q."
+
+ rolled_tensor = tensor.clone()
+
+ cp_size = cp_group.size() if cp_group is not None else 1
+ if cp_size == 1:
+ # CP disabled: simply roll inside each packed sequence boundary.
+ for i in range(len(cu_seqlens) - 1):
+ start_idx = cu_seqlens[i]
+ end_idx = cu_seqlens[i + 1]
+ seq_slice = tensor[..., start_idx:end_idx]
+ rolled_seq = torch.roll(seq_slice, shifts=shifts, dims=dims)
+ rolled_seq[..., shifts:] = 0
+ rolled_tensor[..., start_idx:end_idx] = rolled_seq
+ return rolled_tensor, rolled_tensor.sum()
+
+ # CP enabled: each rank owns two chunks per sequence (front and mirrored tail).
+ local_rank = torch.distributed.get_rank(group=cp_group)
+ global_ranks = torch.distributed.get_process_group_ranks(group=cp_group)
+ next_rank = global_ranks[(local_rank + 1) % cp_size]
+ prev_rank = global_ranks[(local_rank - 1) % cp_size]
+
+ # iterate over each sequence individually
+ for i in range(len(cu_seqlens) - 1):
+ start_idx = cu_seqlens[i]
+ end_idx = cu_seqlens[i + 1]
+
+ # the idx has been multiplied by cp_size, so we need to divide it by cp_size to get the local idx
+ local_start_idx = start_idx // cp_size
+ local_end_idx = end_idx // cp_size
+ tensor_slice = rolled_tensor[..., local_start_idx:local_end_idx].clone()
+
+ # The following code is very similar as the code in roll_tensor function
+ local_chunks = tensor_slice.chunk(2, dim=dims)
+ rolled_chunks = [
+ torch.roll(chunk, shifts=shifts, dims=dims) for chunk in local_chunks
+ ]
+
+ tensor_send_list = []
+ tensor_recv_list = []
+ for chunk in rolled_chunks:
+ boundary = chunk.select(dims, shifts).contiguous().clone()
+ tensor_send_list.append(boundary)
+ tensor_recv_list.append(torch.empty_like(boundary))
+
+ ops = []
+ if local_rank != 0:
+ ops.append(torch.distributed.isend(tensor=tensor_send_list[0], dst=prev_rank))
+ ops.append(torch.distributed.irecv(tensor=tensor_recv_list[1], src=prev_rank))
+ else:
+ tensor_recv_list[1].zero_()
+
+ if local_rank != cp_size - 1:
+ ops.append(torch.distributed.irecv(tensor=tensor_recv_list[0], src=next_rank))
+ ops.append(torch.distributed.isend(tensor=tensor_send_list[1], dst=next_rank))
+ else:
+ tensor_recv_list[0].copy_(tensor_send_list[1])
+
+ for op in ops:
+ op.wait()
+
+ index = [slice(None)] * rolled_chunks[0].dim()
+ index[dims] = shifts
+ for chunk, recv in zip(rolled_chunks, tensor_recv_list):
+ chunk[tuple(index)] = recv
+
+ seq_result = torch.cat(rolled_chunks, dim=dims)
+
+ # update the rolled tensor
+ rolled_tensor[..., local_start_idx:local_end_idx] = seq_result
+
+ return rolled_tensor, rolled_tensor.sum()
class MTPLossLoggingHelper:
"""Helper class for logging MTP losses."""
@@ -480,9 +588,10 @@ class MultiTokenPredictionLayer(MegatronModule):
def _get_embeddings(
self,
input_ids: torch.Tensor,
- position_ids: torch.Tensor,
embedding: Callable,
hidden_states: torch.Tensor,
+ position_ids: Optional[torch.Tensor] = None,
+ packed_seq_params: Optional[PackedSeqParams] = None,
):
"""
Preprocesses input data for the Multi-Token Prediction (MTP) layers.
@@ -499,12 +608,23 @@ class MultiTokenPredictionLayer(MegatronModule):
sequence length, b is the batch size, and h is the hidden size.
"""
# Calc logits for the current Multi-Token Prediction (MTP) layers.
- input_ids, _ = roll_tensor(input_ids, shifts=-1, dims=-1, cp_group=self.cp_group)
- position_ids, _ = roll_tensor(position_ids, shifts=-1, dims=-1, cp_group=self.cp_group)
+ input_ids, _ = roll_tensor(input_ids, shifts=-1, dims=-1, cp_group=self.cp_group, packed_seq_params=packed_seq_params)
+
+ # Prepare/roll position ids only when applicable.
+ if position_ids is None:
+ # Fallback position ids for learned absolute embedding.
+ seq_len = input_ids.size(-1)
+ position_ids = torch.arange(seq_len, dtype=torch.long, device=input_ids.device)
+ position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
+
+ position_ids, _ = roll_tensor(
+ position_ids, shifts=-1, dims=-1, cp_group=self.cp_group, packed_seq_params=packed_seq_params
+ )
# embedding
decoder_input = embedding(input_ids=input_ids, position_ids=position_ids)
+ decoder_input = decoder_input.detach()
- hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True)
+ hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=False)
return input_ids, position_ids, decoder_input, hidden_states
@@ -604,22 +724,66 @@ class MultiTokenPredictionLayer(MegatronModule):
return hidden_states
def _checkpointed_forward(self, forward_func, *args, **kwargs):
+ """Wrap `forward_func` with activation checkpointing while only passing tensors.
+
+ Non-tensor arguments (e.g., configuration objects, None) are captured via closure so
+ that checkpoint implementations never receive them directly, avoiding save_for_backward
+ issues with non-tensor inputs.
+ """
+
+ # TODO(jiajun): Is there any better implementation here?
+ positional_specs = []
+ kw_specs = []
+ tensor_args: List[torch.Tensor] = []
+
+ for arg in args:
+ if torch.is_tensor(arg):
+ positional_specs.append(('tensor', len(tensor_args)))
+ tensor_args.append(arg)
+ else:
+ positional_specs.append(('const', arg))
+
+ for key, value in kwargs.items():
+ if torch.is_tensor(value):
+ kw_specs.append((key, ('tensor', len(tensor_args))))
+ tensor_args.append(value)
+ else:
+ kw_specs.append((key, ('const', value)))
+
+ def run(*flat_tensor_args):
+ rebuilt_args = []
+ for spec_type, payload in positional_specs:
+ if spec_type == 'tensor':
+ rebuilt_args.append(flat_tensor_args[payload])
+ else:
+ rebuilt_args.append(payload)
+
+ rebuilt_kwargs = {}
+ for key, (spec_type, payload) in kw_specs:
+ if spec_type == 'tensor':
+ rebuilt_kwargs[key] = flat_tensor_args[payload]
+ else:
+ rebuilt_kwargs[key] = payload
+
+ return forward_func(*rebuilt_args, **rebuilt_kwargs)
+
+ tensor_args_tuple = tuple(tensor_args)
+
def checkpoint_handler():
- """Determines whether to use the `te_checkpoint` or `tensor_parallel.checkpoint`"""
+ """Determines whether to use the `te_checkpoint` or `tensor_parallel.checkpoint`."""
if self.config.fp8:
from megatron.core.extensions.transformer_engine import te_checkpoint
return te_checkpoint(
- forward_func,
+ run,
self.config.distribute_saved_activations,
tensor_parallel.random.get_cuda_rng_tracker,
parallel_state.get_tensor_model_parallel_group(),
- *args,
- **kwargs,
+ *tensor_args_tuple,
)
else:
return tensor_parallel.checkpoint(
- forward_func, self.config.distribute_saved_activations, *args, *kwargs.values()
+ run, self.config.distribute_saved_activations, *tensor_args_tuple
)
if self.config.recompute_method == 'uniform':
@@ -681,15 +845,13 @@ class MultiTokenPredictionLayer(MegatronModule):
[s, b, h], and optionally the updated context tensor if cross-attention is used.
"""
assert context is None, f"multi token prediction + cross attention is not yet supported."
- assert (
- packed_seq_params is None
- ), f"multi token prediction + sequence packing is not yet supported."
input_ids, position_ids, decoder_input, hidden_states = self._get_embeddings(
input_ids=input_ids,
position_ids=position_ids,
embedding=embedding,
hidden_states=hidden_states,
+ packed_seq_params=packed_seq_params,
)
if self.config.recompute_granularity == 'full' and self.training:
@@ -173,6 +173,10 @@ class TransformerConfig(ModelParallelConfig):
qk_layernorm: bool = False
"""Whether to apply `normalization` type of normalization to the query and key embeddings."""
+ post_self_attn_layernorm: bool = False
+ post_mlp_layernorm: bool = False
+ use_gated_attention: bool = False
+
test_mode: bool = False
"""Whether to run real-time tests."""
@@ -224,6 +224,7 @@ class TransformerLayerSubmodules:
input_layernorm: Union[ModuleSpec, type] = IdentityOp
self_attention: Union[ModuleSpec, type] = IdentityOp
self_attn_bda: Union[ModuleSpec, type] = IdentityFuncOp
+ post_self_attn_layernorm: Union[ModuleSpec, type] = IdentityOp
pre_cross_attn_layernorm: Union[ModuleSpec, type] = IdentityOp
cross_attention: Union[ModuleSpec, type] = IdentityOp
@@ -232,6 +233,7 @@ class TransformerLayerSubmodules:
pre_mlp_layernorm: Union[ModuleSpec, type] = IdentityOp
mlp: Union[ModuleSpec, type] = IdentityOp
mlp_bda: Union[ModuleSpec, type] = IdentityFuncOp
+ post_mlp_layernorm: Union[ModuleSpec, type] = IdentityOp
# Mapping for sharded tensor keys to be applied in `sharded_state_dict` method
sharded_state_dict_keys_map: Dict[str, str] = field(default_factory=dict)
@@ -336,6 +338,13 @@ class TransformerLayer(MegatronModule, BaseTransformerLayer):
# [Module 3: BiasDropoutFusion]
self.self_attn_bda = build_module(submodules.self_attn_bda)
+ self.post_self_attn_layernorm = build_module(
+ submodules.post_self_attn_layernorm,
+ config=self.config,
+ hidden_size=self.config.hidden_size,
+ eps=self.config.layernorm_epsilon,
+ )
+
# [Module 4: Post SelfAttention] Optional Layernorm after self-attn
self.pre_cross_attn_layernorm = build_module(
submodules.pre_cross_attn_layernorm,
@@ -399,6 +408,13 @@ class TransformerLayer(MegatronModule, BaseTransformerLayer):
# [Module 9: BiasDropoutFusion]
self.mlp_bda = build_module(submodules.mlp_bda)
+ self.post_mlp_layernorm = build_module(
+ submodules.post_mlp_layernorm,
+ config=self.config,
+ hidden_size=self.config.hidden_size,
+ eps=self.config.layernorm_epsilon
+ )
+
self.recompute_input_layernorm = False
self.recompute_pre_mlp_layernorm = False
self.recompute_mlp = False
@@ -535,6 +551,10 @@ class TransformerLayer(MegatronModule, BaseTransformerLayer):
attention_output_with_bias[0]
)
+ attention_output, attention_output_bias = attention_output_with_bias
+ attention_output = self.post_self_attn_layernorm(attention_output)
+ attention_output_with_bias = (attention_output, attention_output_bias)
+
# TODO: could we move `bias_dropout_add_exec_handler` itself
# inside the module provided in the `bias_dropout_add_spec` module?
nvtx_range_push(suffix="self_attn_bda")
@@ -635,6 +655,10 @@ class TransformerLayer(MegatronModule, BaseTransformerLayer):
else:
mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output)
+ mlp_output, mlp_output_bias = mlp_output_with_bias
+ mlp_output = self.post_mlp_layernorm(mlp_output)
+ mlp_output_with_bias = (mlp_output, mlp_output_bias)
+
if self.recompute_pre_mlp_layernorm:
# discard the output of the pre-mlp layernorm and register the recompute
# as a gradient hook of mlp_output_with_bias[0]
@@ -937,8 +937,6 @@ def validate_args(args, defaults={}):
# MoE Spec check
if args.num_experts == 0:
args.num_experts = None
- if args.num_experts is not None:
- assert args.spec is None, "Model Spec must be None when using MoEs"
if args.num_experts is not None and args.moe_ffn_hidden_size is None:
args.moe_ffn_hidden_size = args.ffn_hidden_size
print("Warning: moe_ffn_hidden_size is not set, using ffn_hidden_size for MoE instead.")
@@ -1198,6 +1196,10 @@ def core_transformer_config_from_args(args, config_class=None):
if args.is_hybrid_model:
kw_args['is_hybrid_model'] = args.is_hybrid_model
+ kw_args['post_self_attn_layernorm'] = args.post_self_attn_layernorm
+ kw_args['post_mlp_layernorm'] = args.post_mlp_layernorm
+ kw_args['use_gated_attention'] = args.use_gated_attention
+
# handle quantization config
# NOTE: Kitchen arguments are only added to the namespace when
# Kitchen library is available.
@@ -1488,6 +1490,12 @@ def _add_network_size_args(parser):
action='store_true',
help='If set, use original BERT residula connection '
'ordering.')
+ group.add_argument('--post-self-attn-layernorm', action='store_true',
+ help='If set, use post self attention layernorm.')
+ group.add_argument('--post-mlp-layernorm', action='store_true',
+ help='If set, use post MLP layernorm.')
+ group.add_argument('--use-gated-attention', action='store_true',
+ help='If set, use gated attention as in Qwen3Next')
group.add_argument('--openai-gelu', action='store_true',
help='Use OpenAIs GeLU implementation. This option'
'should not be used unless for backward compatibility'
@@ -138,6 +138,8 @@ class _HuggingFaceTokenizer(MegatronTokenizer):
f"The transformers library must be installed to use huggingface_tokenizer_provider"
)
+ if "trust_remote_code" not in kwargs:
+ kwargs["trust_remote_code"] = True
# TODO(bnorick): download tokenizer once to lustre and use force offline to make sure all tasks read it from there
self._tokenizer = transformers.AutoTokenizer.from_pretrained(
pretrained_model_name_or_path=pretrained_model_name_or_path, **kwargs
