diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py index dfa6e4c35e4..2b0a18b433b 100644 --- a/megatron/core/ssm/gated_delta_net.py +++ b/megatron/core/ssm/gated_delta_net.py @@ -21,6 +21,12 @@ from megatron.core.jit import jit_fuser from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.ssm.mamba_context_parallel import ( + _all_to_all_cp2hp, + _all_to_all_hp2cp, + _redo_attention_load_balancing, + _undo_attention_load_balancing, +) from megatron.core.tensor_parallel import get_cuda_rng_tracker from megatron.core.transformer import TransformerConfig from megatron.core.transformer.identity_op import IdentityOp @@ -33,9 +39,6 @@ ) from megatron.core.utils import deprecate_inference_params, nvtx_range_pop, nvtx_range_push -# TODO: Implement GatedDeltaNetContextParallel -# from .gated_delta_net_context_parallel import GatedDeltaNetContextParallel - try: from fla.modules.l2norm import l2norm from fla.ops.gated_delta_rule import chunk_gated_delta_rule @@ -84,6 +87,7 @@ def __init__( use_qk_l2norm: bool = True, A_init_range: Tuple[float, float] = (1, 16), pg_collection: ProcessGroupCollection = None, + **kwargs, ): """ Args: @@ -114,6 +118,7 @@ def __init__( self.use_qk_l2norm = use_qk_l2norm assert pg_collection is not None, "pg_collection must be provided for GatedDeltaNet" self.pg_collection = pg_collection + self.cp_size = self.pg_collection.cp.size() self.tp_size = self.pg_collection.tp.size() self.sp_size = self.tp_size if config.sequence_parallel else 1 @@ -129,6 +134,8 @@ def __init__( self.num_value_heads = config.linear_num_value_heads self.qk_dim = self.key_head_dim * self.num_key_heads self.v_dim = self.value_head_dim * self.num_value_heads + self.qk_dim_local_tp = self.qk_dim // self.tp_size + self.v_dim_local_tp = self.v_dim // self.tp_size # Input projection (hidden_states -> q, k, v, gate, beta, alpha) # TODO: for now, output gate is forced for GDN. @@ -217,8 +224,6 @@ def __init__( tp_group=self.pg_collection.tp, ) - # TODO: support CP - self.reset_parameters() def reset_parameters(self): @@ -247,17 +252,12 @@ def forward( self, hidden_states: Tensor, attention_mask: Tensor, - key_value_states: Optional[Tensor] = None, inference_context: Optional[BaseInferenceContext] = None, - rotary_pos_emb: Optional[Union[Tensor, Tuple[Tensor, Tensor]]] = None, - rotary_pos_cos: Optional[Tensor] = None, - rotary_pos_sin: Optional[Tensor] = None, - rotary_pos_cos_sin: Optional[Tensor] = None, - attention_bias: Optional[Tensor] = None, packed_seq_params: Optional[PackedSeqParams] = None, sequence_len_offset: Optional[int] = None, *, inference_params: Optional[BaseInferenceContext] = None, + **kwargs, ): """ Perform a forward pass through the GDN module. @@ -265,15 +265,8 @@ def forward( Args: hidden_states (Tensor): Hidden states. attention_mask (Tensor): Attention mask. - key_value_states (Optional[Tensor]): Key/value states (for cross attention). inference_context (Optional[BaseInferenceContext]): Inference context that manages KV cache. - rotary_pos_emb (Optional[Union[Tensor, Tuple[Tensor, Tensor]]]): Rotary - embedding tensor(s). - rotary_pos_cos (Optional[Tensor]): Rotary embedding cosine. - rotary_pos_sin (Optional[Tensor]): Rotary embedding sine. - rotary_pos_cos_sin (Optional[Tensor]): Combined rotary embedding cosine and sine. - attention_bias (Optional[Tensor]): Attention bias. packed_seq_params (Optional[PackedSeqparams]): Parameters used for THD format. sequence_len_offset (Optional[int]): Sequence length offset used for inference CUDA graphs. @@ -287,7 +280,7 @@ def forward( inference_context = deprecate_inference_params(inference_context, inference_params) seq_len, batch, _ = hidden_states.shape - seq_len = seq_len * self.sp_size + seq_len = seq_len * self.sp_size * self.cp_size if inference_context is not None: assert ( @@ -306,6 +299,22 @@ def forward( qkvzba, _ = self.in_proj(hidden_states) nvtx_range_pop(suffix="in_proj") + # CP All to All: CP to HP + qkvzba = tensor_a2a_cp2hp( + qkvzba, + seq_dim=0, + head_dim=-1, + cp_group=self.pg_collection.cp, + split_sections=[ + self.qk_dim_local_tp, + self.qk_dim_local_tp, + self.v_dim_local_tp, + self.v_dim_local_tp, + self.num_value_heads // self.tp_size, + self.num_value_heads // self.tp_size, + ], + ) + # Transpose: s b x --> b s x # From sbhd to bshd format qkvzba = qkvzba.transpose(0, 1) @@ -314,10 +323,10 @@ def forward( qkv, gate, beta, alpha = torch.split( qkvzba, [ - (self.qk_dim * 2 + self.v_dim) // self.tp_size, - self.v_dim // self.tp_size, - self.num_value_heads // self.tp_size, - self.num_value_heads // self.tp_size, + (self.qk_dim_local_tp * 2 + self.v_dim_local_tp) // self.cp_size, + self.v_dim_local_tp // self.cp_size, + self.num_value_heads // self.tp_size // self.cp_size, + self.num_value_heads // self.tp_size // self.cp_size, ], dim=-1, ) @@ -328,14 +337,44 @@ def forward( # Convolution on qkv qkv = qkv.transpose(1, 2).contiguous() # b, s, d -> b, d, s nvtx_range_push(suffix="conv1d") + qkv_channels_split_sections = [ + self.qk_dim_local_tp, + self.qk_dim_local_tp, + self.v_dim_local_tp, + ] + conv1d_weight = get_parameter_local_cp( + self.conv1d.weight, + dim=0, + cp_group=self.pg_collection.cp, + split_sections=qkv_channels_split_sections, + ) + conv1d_bias = ( + get_parameter_local_cp( + self.conv1d.bias, + dim=0, + cp_group=self.pg_collection.cp, + split_sections=qkv_channels_split_sections, + ) + if self.conv_bias + else None + ) if (causal_conv1d_fn is None) or self.config.deterministic_mode: - qkv = self.act_fn(self.conv1d(qkv)[..., :seq_len]) + conv_out = F.conv1d( + input=qkv, + weight=conv1d_weight, + bias=conv1d_bias, + stride=self.conv1d.stride, + padding=self.conv1d.padding, + dilation=self.conv1d.dilation, + groups=self.conv_dim_local_tp // self.cp_size, + ) + qkv = self.act_fn(conv_out[..., :seq_len]) else: assert self.activation in ["silu", "swish"] qkv = causal_conv1d_fn( x=qkv, - weight=self.conv1d.weight.squeeze(1), # d, 1, w -> d, w - bias=self.conv1d.bias, + weight=conv1d_weight.squeeze(1), # d, 1, w -> d, w + bias=conv1d_bias, activation=self.activation, ) nvtx_range_pop(suffix="conv1d") @@ -343,7 +382,11 @@ def forward( qkv = qkv.transpose(1, 2) # b, d, s -> b, s, d query, key, value = torch.split( qkv, - [self.qk_dim // self.tp_size, self.qk_dim // self.tp_size, self.v_dim // self.tp_size], + [ + self.qk_dim_local_tp // self.cp_size, + self.qk_dim_local_tp // self.cp_size, + self.v_dim_local_tp // self.cp_size, + ], dim=-1, ) query = query.reshape(batch, seq_len, -1, self.key_head_dim) @@ -367,7 +410,11 @@ def forward( # Calculate g and beta nvtx_range_push(suffix="g_and_beta") - g = -self.A_log.exp() * F.softplus(alpha.float() + self.dt_bias) # In fp32 + A_log_local_cp = get_parameter_local_cp(self.A_log, dim=0, cp_group=self.pg_collection.cp) + dt_bias_local_cp = get_parameter_local_cp( + self.dt_bias, dim=0, cp_group=self.pg_collection.cp + ) + g = -A_log_local_cp.exp() * F.softplus(alpha.float() + dt_bias_local_cp) # In fp32 beta = beta.sigmoid() nvtx_range_pop(suffix="g_and_beta") @@ -406,6 +453,11 @@ def forward( norm_out = norm_out.reshape(batch, seq_len, -1) norm_out = norm_out.transpose(0, 1).contiguous() + # CP all to all: HP to CP + norm_out = tensor_a2a_hp2cp( + norm_out, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp + ) + # Output projection nvtx_range_push(suffix="out_proj") out, out_bias = self.out_proj(norm_out) @@ -479,10 +531,10 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None, tp_gr sharded_state_dict[f"{prefix}in_proj.weight"] = _split_tensor_factory( sharded_state_dict[f"{prefix}in_proj.weight"], [ - self.qk_dim // self.tp_size, - self.qk_dim // self.tp_size, - self.v_dim // self.tp_size, - self.v_dim // self.tp_size, + self.qk_dim_local_tp, + self.qk_dim_local_tp, + self.v_dim_local_tp, + self.v_dim_local_tp, self.num_value_heads // self.tp_size, self.num_value_heads // self.tp_size, ], @@ -502,11 +554,7 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None, tp_gr for conv_layer_name in conv_layer_name_list: sharded_state_dict[f"{prefix}{conv_layer_name}"] = _split_tensor_factory( sharded_state_dict[f"{prefix}{conv_layer_name}"], - [ - self.qk_dim // self.tp_size, - self.qk_dim // self.tp_size, - self.v_dim // self.tp_size, - ], + [self.qk_dim_local_tp, self.qk_dim_local_tp, self.v_dim_local_tp], ["query", "key", "value"], 0, ) @@ -514,6 +562,9 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None, tp_gr return sharded_state_dict +#################### +# Sharded state dict utilities +#################### def _split_tensor_factory( orig_sh_ten: ShardedTensor, split_sections: List[int], split_names: List[str], split_dim: int ) -> ShardedTensorFactory: @@ -574,6 +625,184 @@ def sh_ten_merge_fn(sub_state_dict): ) +#################### +# Context parallel utilities +#################### +def get_parameter_local_cp( + param: torch.Tensor, + dim: int, + cp_group: torch.distributed.ProcessGroup, + split_sections: Optional[List[int]] = None, +) -> torch.Tensor: + """Get the local parameter for the current context parallel rank. + + Args: + param (torch.Tensor): The entire parameter to get the local parameter for. + dim (int): The dimension to split the parameter along. Usually the dimension of head. + cp_group (torch.distributed.ProcessGroup): The context parallel group. + split_sections (Optional[List[int]]): If not None, + first split the parameter along the dimension dim into sections, + then get the local hidden parallel weights separately, + finally concatenate the local hidden parallel weights along the dimension dim. + + Returns: + torch.Tensor: The local parameter for the current context parallel rank. + """ + + cp_size = cp_group.size() + cp_rank = cp_group.rank() + + # No need to split if CP size is 1. + if cp_size == 1: + return param + + # Split first if needed. + if split_sections is not None: + inputs = torch.split(param, split_sections, dim=dim) + outputs = [] + for p in inputs: + p = get_parameter_local_cp(p, dim, cp_group) + outputs.append(p) + return torch.cat(outputs, dim=dim) + + # Slice the parameter. + slices = [slice(None)] * param.dim() + dim_size = param.size(dim=dim) + slices[dim] = slice(cp_rank * dim_size // cp_size, (cp_rank + 1) * dim_size // cp_size) + param = param[slices] + return param + + +def tensor_a2a_cp2hp( + tensor: torch.Tensor, + seq_dim: int, + head_dim: int, + cp_group: torch.distributed.ProcessGroup, + split_sections: Optional[List[int]] = None, + undo_attention_load_balancing: bool = True, +): + """All-to-all context parallel to hidden parallel. + + Args: + tensor (torch.Tensor): The tensor to all-to-all. + Currently only support (seq_len, batch, head_dim) shaped tensor. + seq_dim (int): The dimension of sequence length. Currently only supports seq_dim == 0. + head_dim (int): The dimension of head. Currently only supports head_dim == -1 or 2. + cp_group (torch.distributed.ProcessGroup): The context parallel group. + split_sections (Optional[List[int]]): If not None, split the tensor along the dimension + head_dim into sections first, then do all-to-all for each section separately, + finally concatenate the separated tensors along the dimension head_dim. + undo_attention_load_balancing (bool): Whether to undo the attention load balancing of CP. + + Returns: + torch.Tensor: The all-to-all tensor. + """ + + cp_size = cp_group.size() + + # No need to all-to-all if CP size is 1. + if cp_size == 1: + return tensor + + # Limitations of mamba_context_parallel._all_to_all_cp2hp. + assert seq_dim == 0, f"tensor_a2a_cp2hp only supports seq_dim == 0 for now, but got {seq_dim=}" + assert ( + head_dim == -1 or head_dim == 2 + ), f"tensor_a2a_cp2hp only supports head_dim == -1 or 2 for now, but got {head_dim=}" + assert ( + tensor.dim() == 3 + ), f"tensor_a2a_cp2hp only supports 3-d input tensor for now, but got {tensor.dim()=}" + + # Split first if needed. + if split_sections is not None: + inputs = torch.split(tensor, split_sections, dim=head_dim) + outputs = [] + for x in inputs: + x = tensor_a2a_cp2hp( + x, + seq_dim=seq_dim, + head_dim=head_dim, + cp_group=cp_group, + undo_attention_load_balancing=False, + ) + outputs.append(x) + tensor = torch.cat(outputs, dim=head_dim) + else: + tensor = _all_to_all_cp2hp(tensor, cp_group) + + # Undo attention load balancing last if needed. + if undo_attention_load_balancing: + tensor = _undo_attention_load_balancing(tensor, cp_size) + return tensor + + +def tensor_a2a_hp2cp( + tensor: torch.Tensor, + seq_dim: int, + head_dim: int, + cp_group: torch.distributed.ProcessGroup, + split_sections: Optional[List[int]] = None, + redo_attention_load_balancing: bool = True, +): + """All-to-all hidden parallel to context parallel. + + Args: + tensor (torch.Tensor): The tensor to all-to-all. + Currently only support (seq_len, batch, head_dim) shaped tensor. + seq_dim (int): The dimension of sequence length. Currently only supports seq_dim == 0. + head_dim (int): The dimension of head. Currently only supports head_dim == -1 or 2. + cp_group (torch.distributed.ProcessGroup): The context parallel group. + split_sections (Optional[List[int]]): If not None, first split the tensor along the + dimension head_dim into sections, then do all-to-all for each section separately, + finally concatenate the separated tensors along the dimension head_dim. + redo_attention_load_balancing (bool): Whether to redo the attention load balancing of HP. + + Returns: + torch.Tensor: The all-to-all tensor. + """ + + cp_size = cp_group.size() + + # No need to all-to-all if CP size is 1. + if cp_size == 1: + return tensor + + # Limitations of mamba_context_parallel._all_to_all_hp2cp. + assert seq_dim == 0, f"tensor_a2a_cp2hp only supports seq_dim == 0 for now, but got {seq_dim=}" + assert ( + head_dim == -1 or head_dim == 2 + ), f"tensor_a2a_cp2hp only supports head_dim == -1 or 2 for now, but got {head_dim=}" + assert ( + tensor.dim() == 3 + ), f"tensor_a2a_cp2hp only supports 3-d input tensor for now, but got {tensor.dim()=}" + + # Redo attention load balancing first if needed. + if redo_attention_load_balancing: + tensor = _redo_attention_load_balancing(tensor, cp_size) + + # Split first if needed. + if split_sections is not None: + inputs = torch.split(tensor, split_sections, dim=head_dim) + outputs = [] + for x in inputs: + x = tensor_a2a_hp2cp( + x, + seq_dim=seq_dim, + head_dim=head_dim, + cp_group=cp_group, + redo_attention_load_balancing=False, + ) + outputs.append(x) + tensor = torch.cat(outputs, dim=head_dim) + else: + tensor = _all_to_all_hp2cp(tensor, cp_group) + + return tensor + + +#################### +# Torch native gated delta rule +#################### def torch_chunk_gated_delta_rule( query, key, diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index e2705bd9f51..6493a4bcce1 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -922,17 +922,14 @@ def __post_init__(self): ) # Check tensor parallelism compatibility - assert ( - self.linear_num_key_heads % self.tensor_model_parallel_size == 0 - ), "linear_num_key_heads must be a multiple of tensor_model_parallel_size." - assert ( - self.linear_num_value_heads % self.tensor_model_parallel_size == 0 - ), "linear_num_value_heads must be a multiple of tensor_model_parallel_size." - - # Do not support yet, but coming soon. - assert self.context_parallel_size == 1, ( - f"Gated delta net does not support context parallel for now," - f" but got {self.context_parallel_size=}." + tp_cp_size = self.tensor_model_parallel_size * self.context_parallel_size + assert self.linear_num_key_heads % tp_cp_size == 0, ( + f"{self.linear_num_key_heads=} must be a multiple of " + f"({self.tensor_model_parallel_size=} * {self.context_parallel_size=})." + ) + assert self.linear_num_value_heads % tp_cp_size == 0, ( + f"{self.linear_num_value_heads=} must be a multiple of " + f"({self.tensor_model_parallel_size=} * {self.context_parallel_size=})." ) elif self.experimental_attention_variant == "dsa": assert ( diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py index 89a185e3755..725d18fbc06 100644 --- a/tests/unit_tests/ssm/test_gated_delta_net.py +++ b/tests/unit_tests/ssm/test_gated_delta_net.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from functools import partial from unittest import mock @@ -28,6 +28,7 @@ init_checkpointing_mock_args, ) from tests.unit_tests.test_utilities import Utils +from tests.unit_tests.transformer.test_attention import _test_parallel_attention_correctness try: import fla @@ -39,12 +40,7 @@ @pytest.mark.parametrize( ("tp_size", "sp", "cp_size"), - [ - (1, False, 1), - (2, False, 1), - (2, True, 1), - # GDN does not support CP for now. Leave it for future work. - ], + [(1, False, 1), (2, False, 1), (2, True, 1), (1, False, 2), (2, False, 2), (2, True, 2)], ) @pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") @pytest.mark.internal @@ -142,50 +138,13 @@ def test_gpu_forward(self): [ (4, False, 1), # TP w/o SP (4, True, 1), # TP w/ SP - # CP does not support GDN for now. Add it once it is supported. + (1, False, 2), # CP + (2, False, 2), # TP w/o SP + CP + (2, True, 2), # TP w/ SP + CP ], ) @pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") def test_parallel_gated_delta_net_correctness(tmp_path_dist_ckpt, tp, sp, cp): - # Constants - seed = 123 - sequence_length = 256 - micro_batch_size = 4 - hidden_size = 128 - normalization = "RMSNorm" - - # Model initialization function - def initialize_gpt_model(config, pre_process=True, post_process=True, vp_stage=None): - layer_spec = get_gpt_layer_with_transformer_engine_spec( - experimental_attention_variant="gated_delta_net", normalization=normalization - ) - gpt_model = GPTModel( - config=config, - transformer_layer_spec=layer_spec, - vocab_size=128, - max_sequence_length=sequence_length, - pre_process=pre_process, - post_process=post_process, - vp_stage=vp_stage, - ) - return gpt_model - - # Initialize baseline parallel state - Utils.initialize_model_parallel( - tensor_model_parallel_size=1, pipeline_model_parallel_size=1, context_parallel_size=1 - ) - - # Initialize input hidden states - torch.manual_seed(seed) - model_parallel_cuda_manual_seed(seed) - input_hidden_states = ( - torch.rand((sequence_length, micro_batch_size, hidden_size)) - .cuda() - .bfloat16() - .requires_grad_(True) - ) - - # Initialize transformer config transformer_config = TransformerConfig( hidden_size=128, linear_conv_kernel_dim=2, @@ -194,7 +153,7 @@ def initialize_gpt_model(config, pre_process=True, post_process=True, vp_stage=N linear_num_key_heads=4, linear_num_value_heads=8, num_layers=1, - normalization=normalization, + normalization="RMSNorm", use_cpu_initialization=True, layernorm_zero_centered_gamma=True, num_attention_heads=8, @@ -202,118 +161,15 @@ def initialize_gpt_model(config, pre_process=True, post_process=True, vp_stage=N bf16=True, ) - with TempNamedDir(tmp_path_dist_ckpt / 'test_parallel_gdn', sync=True) as ckpt_dir: - # Set argument - mock_args = parse_args(ignore_unknown_args=True) - set_args(mock_args) - - # Initialize baseline model - init_basic_mock_args(mock_args, 1, 1, bf16=True) - mock_args.context_parallel_size = 1 - mock_args.sequence_parallel = 1 - gpt_model = unwrap_model( - get_model(partial(initialize_gpt_model, config=transformer_config)) - ) - - # Initialize args and save checkpoint - init_checkpointing_mock_args(mock_args, ckpt_dir, False) - mock_args.no_save_optim = True - mock_args.no_save_rng = True - mock_args.no_load_optim = True - mock_args.no_load_rng = True - save_checkpoint(10, gpt_model, None, None, 0) - - # Calculate baseline output - attention = gpt_model[0].decoder.layers[0].self_attention - output_hidden_states_baseline, bias_hidden_states_baseline = attention( - input_hidden_states, attention_mask=None - ) - output_hidden_states_baseline.sum().backward() - - # Save baseline output - input_grad_baseline = input_hidden_states.grad.detach() - output_hidden_states_baseline = output_hidden_states_baseline.detach() - - # Initialize parallel model - Utils.destroy_model_parallel() - Utils.initialize_model_parallel( - tensor_model_parallel_size=tp, pipeline_model_parallel_size=1, context_parallel_size=cp - ) - torch.manual_seed(seed) - model_parallel_cuda_manual_seed(seed) - transformer_config.context_parallel_size = cp - transformer_config.tensor_model_parallel_size = tp - transformer_config.sequence_parallel = sp - init_basic_mock_args(mock_args, tp, 1, bf16=True) - mock_args.context_parallel_size = cp - mock_args.sequence_parallel = sp - gpt_model = unwrap_model( - get_model(partial(initialize_gpt_model, config=transformer_config)) - ) - with mock.patch('megatron.training.checkpointing.check_checkpoint_args'): - with mock.patch('megatron.training.checkpointing.update_num_microbatches'): - load_checkpoint(gpt_model, None, None) - - # Function to get tensor on this tp and cp rank - cp_group = parallel_state.get_context_parallel_group() - tp_rank = parallel_state.get_tensor_model_parallel_rank() - - def get_tensor_on_this_rank(tensor): - if cp > 1: - tensor = get_tensor_on_this_cp_rank(tensor, 0, cp_group) - if tp > 1 and sp: - sp_seg = sequence_length // tp // cp - tensor = tensor[tp_rank * sp_seg : (tp_rank + 1) * sp_seg] - return tensor - - # Calculate parallel model output - input_hidden_states = get_tensor_on_this_rank(input_hidden_states) - input_hidden_states = input_hidden_states.detach().requires_grad_(True) - parallel_attention = gpt_model[0].decoder.layers[0].self_attention - output_hidden_states_parallel, bias_hidden_states_parallel = parallel_attention( - input_hidden_states, attention_mask=None - ) - output_hidden_states_parallel.sum().backward() - input_grad_parallel = input_hidden_states.grad.detach() - - # Check if the output is the same - if cp: - atol, rtol = 5e-3, 5e-3 - else: - atol, rtol = 5e-4, 5e-4 - output_hidden_states_baseline = get_tensor_on_this_rank(output_hidden_states_baseline) - input_grad_baseline = get_tensor_on_this_rank(input_grad_baseline) - - assert torch.all( - ~torch.isnan(output_hidden_states_baseline) - ), "output_hidden_states_baseline contains nan" - assert torch.all( - ~torch.isinf(output_hidden_states_baseline) - ), "output_hidden_states_baseline contains inf" - assert torch.all(~torch.isnan(input_grad_baseline)), "input_grad_baseline contains nan" - assert torch.all(~torch.isinf(input_grad_baseline)), "input_grad_baseline contains inf" - assert torch.all( - ~torch.isnan(output_hidden_states_parallel) - ), "output_hidden_states_parallel contains nan" - assert torch.all( - ~torch.isinf(output_hidden_states_parallel) - ), "output_hidden_states_parallel contains inf" - assert torch.all(~torch.isnan(input_grad_parallel)), "input_grad_parallel contains nan" - assert torch.all(~torch.isinf(input_grad_parallel)), "input_grad_parallel contains inf" + transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec( + experimental_attention_variant="gated_delta_net", normalization="RMSNorm" + ) - torch.testing.assert_close( - output_hidden_states_baseline, - output_hidden_states_parallel, - atol=atol, - rtol=rtol, - msg=lambda msg: f"Mismatch in output_hidden_states: {msg}", - ) - torch.testing.assert_close( - input_grad_baseline, - input_grad_parallel, - atol=atol, - rtol=rtol, - msg=lambda msg: f"Mismatch in input_grad: {msg}", - ) + if cp: + atol, rtol = 5e-3, 5e-3 + else: + atol, rtol = 5e-4, 5e-4 - Utils.destroy_model_parallel() + _test_parallel_attention_correctness( + transformer_config, transformer_layer_spec, tmp_path_dist_ckpt, tp, sp, cp + )