diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 54981c9086..497f414aee 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -6,6 +6,11 @@ PyTorch ======= +.. autoapiclass:: transformer_engine.pytorch.autocast(enabled=True, calibrating=False, recipe=None, amax_reduction_group=None) + +Standard layers +--------------- + .. autoapiclass:: transformer_engine.pytorch.Linear(in_features, out_features, bias=True, **kwargs) :members: forward, set_tensor_parallel_group @@ -34,20 +39,34 @@ PyTorch .. autoapiclass:: transformer_engine.pytorch.TransformerLayer(hidden_size, ffn_hidden_size, num_attention_heads, **kwargs) :members: forward, set_context_parallel_group, set_tensor_parallel_group +Model-specific layers +--------------------- + +DeepSeek-V3 +^^^^^^^^^^^ + +.. autoapiclass:: transformer_engine.pytorch.models.DeepSeekV3Layer(hidden_size, num_attention_heads, **kwargs) + :members: forward + +.. autoapiclass:: transformer_engine.pytorch.models.DeepSeekV3MoE(hidden_size, moe_ffn_hidden_size, num_experts, **kwargs) + :members: forward, update_expert_bias + +.. autoapiclass:: transformer_engine.pytorch.models.MultiLatentAttention(hidden_size, num_attention_heads, **kwargs) + :members: forward + +Other +----- + .. autoapiclass:: transformer_engine.pytorch.dot_product_attention.inference.InferenceParams(max_batch_size, max_sequence_length) :members: reset, allocate_memory, pre_step, get_seqlens_pre_step, convert_paged_to_nonpaged, step .. autoapiclass:: transformer_engine.pytorch.CudaRNGStatesTracker() :members: reset, get_states, set_states, add, fork - -.. autoapiclass:: transformer_engine.pytorch.autocast(enabled=True, calibrating=False, recipe=None, amax_reduction_group=None) - .. autoapifunction:: transformer_engine.pytorch.quantized_model_init .. autoapifunction:: transformer_engine.pytorch.checkpoint - .. autoapifunction:: transformer_engine.pytorch.make_graphed_callables .. autoapifunction:: transformer_engine.pytorch.get_cpu_offload_context diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index a78a99d7f9..56aef3b840 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -58,6 +58,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_backward_overrid python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_permutation.xml $TE_PATH/tests/pytorch/test_permutation.py || test_fail "test_permutation.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cross_entropy.xml $TE_PATH/tests/pytorch/test_cross_entropy.py || test_fail "test_cross_entropy.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading.xml $TE_PATH/tests/pytorch/test_cpu_offloading.py || test_fail "test_cpu_offloading.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_models.xml $TE_PATH/tests/pytorch/test_models.py || test_fail "test_models.py" NVTE_FLASH_ATTN=0 NVTE_CPU_OFFLOAD_V1=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading_v1.xml $TE_PATH/tests/pytorch/test_cpu_offloading_v1.py || test_fail "test_cpu_offloading_v1.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hybrid_quantization.xml $TE_PATH/tests/pytorch/test_hybrid_quantization.py || test_fail "test_hybrid_quantization.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_identity_quantizer.xml $TE_PATH/tests/pytorch/test_identity_quantizer.py || test_fail "test_identity_quantizer.py" diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh index f1de313fdc..68f242870c 100644 --- a/qa/L1_pytorch_distributed_unittest/test.sh +++ b/qa/L1_pytorch_distributed_unittest/test.sh @@ -56,6 +56,7 @@ python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_cu_seqlens_cache.xml python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_cast_master_weights_to_fp8.xml $TE_PATH/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py || test_fail "test_cast_master_weights_to_fp8.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_newton_schulz.xml $TE_PATH/tests/pytorch/distributed/test_newton_schulz.py || test_fail "test_newton_schulz.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_ep.xml $TE_PATH/tests/pytorch/distributed/test_ep.py || test_fail "test_ep.py" +python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_models.xml $TE_PATH/tests/pytorch/distributed/test_models.py || test_fail "distributed/test_models.py" # debug tests diff --git a/tests/pytorch/attention/test_linear_mxfp8_attention.py b/tests/pytorch/attention/test_linear_mxfp8_attention.py index f1bba7bc9a..95770a6bb8 100644 --- a/tests/pytorch/attention/test_linear_mxfp8_attention.py +++ b/tests/pytorch/attention/test_linear_mxfp8_attention.py @@ -36,7 +36,11 @@ _current_file = pathlib.Path(__file__).resolve() sys.path = [str(_current_file.parent.parent)] + sys.path from utils import ModelConfig, compare_and_assert, get_available_attention_backends -from mla_rope_utils import apply_mla_rope, build_rope_tables +from transformer_engine.pytorch.models.deepseek_v3.mla_rope import ( + apply_mla_rope_kv, + apply_mla_rope_q, + build_rope_tables, +) try: @@ -183,6 +187,13 @@ def _run_projections( return q_flat, kv_flat, q, kv, k_pos_emb +def _apply_rope(q, kv, k_pos_emb, rope_tables): + cos, sin = rope_tables + q = apply_mla_rope_q(q, cos, sin, HEAD_DIM_NOPE, HEAD_DIM_ROPE) + k, v = apply_mla_rope_kv(kv, k_pos_emb, cos, sin, HEAD_DIM_NOPE, HEAD_DIM_ROPE, HEAD_DIM_V) + return q, k, v + + def _run_forward_bf16( modules: tuple, x: torch.Tensor, @@ -190,7 +201,7 @@ def _run_forward_bf16( ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: q_proj, kv_proj, dpa, out_linear = modules _, _, q, kv, k_pos_emb = _run_projections(q_proj, kv_proj, x) - q, k, v = apply_mla_rope(q, kv, k_pos_emb, cos_table=rope_tables[0], sin_table=rope_tables[1]) + q, k, v = _apply_rope(q, kv, k_pos_emb, rope_tables) attn_out = dpa(q, k, v, qkv_format="sbhd") return q, k, v, out_linear(attn_out.view(x.shape[0], x.shape[1], HIDDEN_SIZE)) @@ -212,13 +223,7 @@ def _run_forward_mxfp8( x, is_first_microbatch, ) - q, k, v = apply_mla_rope( - q, - kv, - k_pos_emb, - cos_table=rope_tables[0], - sin_table=rope_tables[1], - ) + q, k, v = _apply_rope(q, kv, k_pos_emb, rope_tables) attn_out = dpa(q, k, v, qkv_format="sbhd") out = out_linear( attn_out.view(x.shape[0], x.shape[1], HIDDEN_SIZE), @@ -292,7 +297,7 @@ def test_accuracy(self, batch_size: int, seq_len: int) -> None: _set_seed() baseline_modules, mxfp8_modules = _build_modules() x = torch.randn(seq_len, batch_size, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") - rope_tables = build_rope_tables(seq_len, device=x.device) + rope_tables = build_rope_tables(seq_len, HEAD_DIM_ROPE, device=x.device) q_bf16, k_bf16, v_bf16, out_bf16 = _run_forward_bf16(baseline_modules, x, rope_tables) q_mxfp8, k_mxfp8, v_mxfp8, out_mxfp8 = _run_forward_mxfp8( @@ -378,7 +383,7 @@ def test_backward(self, batch_size: int, seq_len: int) -> None: device="cuda", requires_grad=True, ) - rope_tables = build_rope_tables(seq_len, device=x.device) + rope_tables = build_rope_tables(seq_len, HEAD_DIM_ROPE, device=x.device) *_, out_mxfp8 = _run_forward_mxfp8(mxfp8_modules, x, fp8_recipe, rope_tables) out_mxfp8.sum().backward() @@ -412,7 +417,7 @@ def test_performance(self, batch_size: int, seq_len: int) -> None: device="cuda", requires_grad=True, ) - rope_tables = build_rope_tables(seq_len, device=x.device) + rope_tables = build_rope_tables(seq_len, HEAD_DIM_ROPE, device=x.device) mxfp8_fprop_ms, mxfp8_bprop_ms = _benchmark_training_step( _run_forward_mxfp8, mxfp8_modules, x, fp8_recipe, rope_tables diff --git a/tests/pytorch/distributed/run_models.py b/tests/pytorch/distributed/run_models.py new file mode 100644 index 0000000000..9561d20117 --- /dev/null +++ b/tests/pytorch/distributed/run_models.py @@ -0,0 +1,169 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Multi-process tests for model-specific layers (te.models), launched via torchrun.""" + +import os +import sys + +import torch +import torch.distributed as dist + +from transformer_engine.pytorch.ep import ep_bootstrap, ep_finalize, release_symm_mem_pool +from transformer_engine.pytorch.models import DeepSeekV3Layer + +HIDDEN = 256 +MOE_FFN = 128 +SHARED_FFN = 128 +NUM_LOCAL_EXPERTS = 2 +TOP_K = 2 +TOKENS_PER_RANK = 64 +HEADS = 4 +DTYPE = torch.bfloat16 + +MLA_KWARGS = dict( + q_lora_rank=96, + kv_lora_rank=64, + qk_nope_head_dim=64, + qk_rope_head_dim=32, + v_head_dim=64, +) + + +def _device_sm() -> int: + major, minor = torch.cuda.get_device_capability() + return major * 10 + minor + + +def _recv_capacity(ep_size: int) -> int: + cap = ep_size * TOKENS_PER_RANK * TOP_K + NUM_LOCAL_EXPERTS * 128 + return -(-cap // 128) * 128 + + +def _broadcast_params(module: torch.nn.Module) -> None: + for t in list(module.parameters()) + list(module.buffers()): + dist.broadcast(t.detach(), src=0) + + +def _make_layer(ep_group, ep_size: int, num_experts: int) -> DeepSeekV3Layer: + ep = ep_group is not None + return DeepSeekV3Layer( + HIDDEN, + HEADS, + num_experts=num_experts, + moe_ffn_hidden_size=MOE_FFN, + topk=TOP_K, + shared_expert_ffn_hidden_size=SHARED_FFN, + params_dtype=DTYPE, + ep_group=ep_group, + ep_max_tokens_per_rank=TOKENS_PER_RANK if ep else None, + **MLA_KWARGS, + ) + + +def _copy_weights(ep_layer: DeepSeekV3Layer, ref: DeepSeekV3Layer, rank: int) -> None: + ref_params = dict(ref.named_parameters()) + ref_bufs = dict(ref.named_buffers()) + with torch.no_grad(): + for name, p in ep_layer.named_parameters(): + if not name.startswith("mlp.experts."): + p.copy_(ref_params[name]) + for name, b in ep_layer.named_buffers(): + if name in ref_bufs and b.shape == ref_bufs[name].shape: + b.copy_(ref_bufs[name]) + ep_fc1, _, ep_fc2 = ep_layer.mlp.experts + ref_fc1, _, ref_fc2 = ref.mlp.experts + for local_e in range(NUM_LOCAL_EXPERTS): + global_e = rank * NUM_LOCAL_EXPERTS + local_e + getattr(ep_fc1, f"weight{local_e}").copy_(getattr(ref_fc1, f"weight{global_e}")) + getattr(ep_fc2, f"weight{local_e}").copy_(getattr(ref_fc2, f"weight{global_e}")) + + +def test_layer_ep_matches_local(rank: int, ep_size: int, ep_group) -> None: + """Full DeepSeekV3Layer with EP must match the all-experts-local layer numerically.""" + num_experts = NUM_LOCAL_EXPERTS * ep_size + torch.manual_seed(0) + ref = _make_layer(None, ep_size, num_experts) + _broadcast_params(ref) + ep_layer = _make_layer(ep_group, ep_size, num_experts) + _copy_weights(ep_layer, ref, rank) + + torch.manual_seed(1234 + rank) + x = torch.randn(TOKENS_PER_RANK // 2, 2, HIDDEN, dtype=DTYPE, device="cuda") + x_ep = x.clone().requires_grad_(True) + x_ref = x.clone().requires_grad_(True) + + out_ep = ep_layer(x_ep) + out_ref = ref(x_ref) + assert out_ep.shape == x.shape + torch.testing.assert_close(out_ep, out_ref, rtol=0.05, atol=0.05) + + grad_out = torch.randn_like(out_ep) + out_ep.backward(grad_out) + out_ref.backward(grad_out) + torch.testing.assert_close(x_ep.grad, x_ref.grad, rtol=0.05, atol=0.05) + + ref_params = dict(ref.named_parameters()) + for name, p in ep_layer.named_parameters(): + if name.startswith("mlp.experts.") or p.grad is None: + continue + torch.testing.assert_close(p.grad, ref_params[name].grad, rtol=0.1, atol=0.1, msg=name) + + # A local expert's wgrad on its owner rank equals the sum of the + # reference wgrads over all ranks. all_reduce is collective, so every + # rank must reduce every expert's grad (in the same order). + ep_fc1, _, ep_fc2 = ep_layer.mlp.experts + ref_fc1, _, ref_fc2 = ref.mlp.experts + for ep_fc, ref_fc in ((ep_fc1, ref_fc1), (ep_fc2, ref_fc2)): + ref_grads = [getattr(ref_fc, f"weight{e}").grad.float().clone() for e in range(num_experts)] + for g in ref_grads: + dist.all_reduce(g) + for local_e in range(NUM_LOCAL_EXPERTS): + global_e = rank * NUM_LOCAL_EXPERTS + local_e + ep_grad = getattr(ep_fc, f"weight{local_e}").grad.float() + torch.testing.assert_close(ep_grad, ref_grads[global_e], rtol=0.1, atol=0.1) + + counts = ep_layer.mlp._last_tokens_per_expert.clone() + dist.all_reduce(counts) + assert counts.sum().item() == ep_size * TOKENS_PER_RANK * TOP_K + + ep_layer.mlp.update_expert_bias() + assert torch.isfinite(ep_layer.mlp.expert_bias).all() + + +def main() -> int: + dist.init_process_group(backend="nccl") + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + from torch.distributed import _symmetric_memory as _symm_mem + + _symm_mem.set_backend("NCCL") + + rank = dist.get_rank() + ep_size = dist.get_world_size() + if _device_sm() < 90: + if rank == 0: + print(f"NCCL EP requires SM>=90 (got SM{_device_sm()}); skipping.") + dist.destroy_process_group() + return 0 + + ep_group = dist.new_group(ranks=list(range(ep_size)), backend="nccl") + ep_bootstrap( + ep_group, + num_experts=NUM_LOCAL_EXPERTS * ep_size, + max_tokens_per_rank=TOKENS_PER_RANK, + hidden_dim=HIDDEN, + num_topk=TOP_K, + recv_capacity_per_rank=_recv_capacity(ep_size), + ) + test_layer_ep_matches_local(rank, ep_size, ep_group) + print(f"[rank {rank}] PASSED") + + dist.barrier() + ep_finalize() + release_symm_mem_pool() + dist.destroy_process_group() + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/pytorch/distributed/test_models.py b/tests/pytorch/distributed/test_models.py new file mode 100644 index 0000000000..1b96eae2aa --- /dev/null +++ b/tests/pytorch/distributed/test_models.py @@ -0,0 +1,31 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import os +import subprocess +from pathlib import Path + +import pytest +import torch + +TEST_ROOT = Path(__file__).parent.resolve() +NUM_PROCS = min(8, torch.cuda.device_count()) +LAUNCH_CMD = ["torchrun", f"--nproc_per_node={NUM_PROCS}"] + + +def _has_nvlink() -> bool: + # NCCL EP falls back to the network transport and deadlocks on PCIe-only nodes. + out = subprocess.run( + ["nvidia-smi", "nvlink", "--status"], capture_output=True, text=True, check=False + ).stdout + return "GB/s" in out + + +@pytest.mark.skipif(NUM_PROCS < 2, reason="EP requires >= 2 GPUs") +@pytest.mark.skipif(not _has_nvlink(), reason="NCCL EP requires NVLink") +def test_deepseek_layer_ep(): + result = subprocess.run( + LAUNCH_CMD + [str(TEST_ROOT / "run_models.py")], env=os.environ, check=False, timeout=300 + ) + assert result.returncode == 0 diff --git a/tests/pytorch/test_models.py b/tests/pytorch/test_models.py new file mode 100644 index 0000000000..cd1903e79f --- /dev/null +++ b/tests/pytorch/test_models.py @@ -0,0 +1,176 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import math + +import pytest +import torch + +from transformer_engine.pytorch.utils import deinterleave_glu_tensor +from transformer_engine.pytorch.models import DeepSeekV3MoE, MultiLatentAttention + +SEQ_LEN = 128 +BATCH = 2 +HIDDEN = 256 +HEADS = 4 +DTYPE = torch.bfloat16 + +MLA_KWARGS = dict( + q_lora_rank=96, + kv_lora_rank=64, + qk_nope_head_dim=64, + qk_rope_head_dim=32, + v_head_dim=64, +) + + +def _input(requires_grad=True): + torch.manual_seed(1234) + return torch.randn( + SEQ_LEN, BATCH, HIDDEN, dtype=DTYPE, device="cuda", requires_grad=requires_grad + ) + + +def test_mla_rope_triton_matches_pytorch(): + from transformer_engine.pytorch.models.deepseek_v3 import mla_rope + + if not mla_rope.HAVE_TRITON: + pytest.skip("Triton unavailable") + s, b, h = 64, 2, 4 + nope, rope, vdim = 64, 32, 64 + cos, sin = mla_rope.build_rope_tables(s, rope, device="cuda") + + torch.manual_seed(0) + q_leaf = torch.randn(s, b, h, nope + rope, device="cuda", requires_grad=True) + kv_leaf = torch.randn(s, b, h, nope + vdim, device="cuda", requires_grad=True) + pos_leaf = torch.randn(s, b, 1, rope, device="cuda", requires_grad=True) + grad_q = torch.randn(s, b, h, nope + rope, device="cuda") + grad_k = torch.randn(s, b, h, nope + rope, device="cuda") + grad_v = torch.randn(s, b, h, vdim, device="cuda") + + def run(fmt): + # non-leaf copies: the Triton q kernel rotates in place + q, kv, pos = q_leaf * 1.0, kv_leaf * 1.0, pos_leaf * 1.0 + q_out = mla_rope.apply_mla_rope_q(q, cos, sin, nope, rope, fmt) + k_out, v_out = mla_rope.apply_mla_rope_kv(kv, pos, cos, sin, nope, rope, vdim, fmt) + # fresh grad clones: the Triton q backward modifies its input grad in place + torch.autograd.backward( + [q_out, k_out, v_out], [grad_q.clone(), grad_k.clone(), grad_v.clone()] + ) + grads = (q_leaf.grad.clone(), kv_leaf.grad.clone(), pos_leaf.grad.clone()) + q_leaf.grad = kv_leaf.grad = pos_leaf.grad = None + return (q_out.clone(), k_out, v_out), grads + + (q_t, k_t, v_t), grads_t = run("sbhd") + + seq_dim = 0 + q_ref = torch.cat( + ( + (q_leaf * 1.0)[..., :nope], + mla_rope._rotate_interleaved_to_neox((q_leaf * 1.0)[..., nope:], cos, sin, seq_dim), + ), + dim=-1, + ) + k_ref = torch.cat( + ( + (kv_leaf * 1.0)[..., :nope], + mla_rope._rotate_interleaved_to_neox(pos_leaf * 1.0, cos, sin, seq_dim).expand( + s, b, h, rope + ), + ), + dim=-1, + ) + v_ref = (kv_leaf * 1.0)[..., nope:] + torch.autograd.backward([q_ref, k_ref, v_ref], [grad_q.clone(), grad_k.clone(), grad_v.clone()]) + + torch.testing.assert_close(q_t, q_ref, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(k_t, k_ref, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(v_t, v_ref, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(grads_t[0], q_leaf.grad, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(grads_t[1], kv_leaf.grad, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(grads_t[2], pos_leaf.grad, rtol=1e-5, atol=1e-5) + + +def test_rope_tables_yarn(): + from transformer_engine.pytorch.models.deepseek_v3 import mla_rope + + s, rope = 8192, 64 + cos, sin = mla_rope.build_rope_tables(s, rope, device="cuda") + cos_none, sin_none = mla_rope.build_rope_tables(s, rope, device="cuda", scaling_factor=None) + assert torch.equal(cos, cos_none) and torch.equal(sin, sin_none) + + yarn = dict(scaling_factor=40.0, original_max_position_embeddings=4096) + cos_y, sin_y = mla_rope.build_rope_tables(s, rope, device="cuda", **yarn) + factor = mla_rope.yarn_concentration_factor(40.0, 1.0, 0.0) + assert factor == pytest.approx(0.1 * math.log(40.0) + 1.0) + # amplitude scaled by the concentration factor + torch.testing.assert_close(cos_y**2 + sin_y**2, torch.full_like(cos_y, factor**2)) + # high-frequency dims untouched, low-frequency dims interpolated by 1/scaling_factor + torch.testing.assert_close(cos_y[:, 0] / factor, cos[:, 0]) + angle_y = torch.atan2(sin_y[:, rope // 2 - 1], cos_y[:, rope // 2 - 1]) + angle = torch.atan2(sin[:, rope // 2 - 1], cos[:, rope // 2 - 1]) + torch.testing.assert_close(angle_y[:64], angle[:64] / 40.0, atol=1e-4, rtol=0) + + +@pytest.mark.parametrize("mscale_all_dim", [0.0, 1.0]) +def test_mla_yarn_softmax_scale(mscale_all_dim): + mla = MultiLatentAttention( + HIDDEN, + HEADS, + params_dtype=DTYPE, + rope_scaling_factor=40.0, + original_max_position_embeddings=64, + mscale_all_dim=mscale_all_dim, + **MLA_KWARGS, + ) + m = 0.1 * mscale_all_dim * math.log(40.0) + 1.0 + qk_head_dim = MLA_KWARGS["qk_nope_head_dim"] + MLA_KWARGS["qk_rope_head_dim"] + assert mla.softmax_scale == pytest.approx(m * m / math.sqrt(qk_head_dim)) + + +@pytest.mark.parametrize("shared", [False, True], ids=["no_shared", "shared"]) +@pytest.mark.parametrize("grouped", [False, True], ids=["ungrouped", "grouped"]) +@pytest.mark.parametrize("topk", [2, 4]) +def test_moe_matches_dense_reference(shared, grouped, topk): + """Routed output must equal the prob-weighted sum of the selected expert MLPs.""" + torch.manual_seed(0) + num_experts = 4 + moe = DeepSeekV3MoE( + HIDDEN, + moe_ffn_hidden_size=128, + num_experts=num_experts, + topk=topk, + num_groups=2 if grouped else None, + group_topk=topk // 2 if grouped else None, + shared_expert_ffn_hidden_size=128 if shared else None, + params_dtype=DTYPE, + ) + x = _input() + out = moe(x) + assert out.shape == x.shape + out.sum().backward() + assert torch.isfinite(x.grad).all() + + tokens = x.detach().reshape(-1, HIDDEN) + probs, _ = moe._route(moe.gate(tokens).float()) + assert (probs > 0).sum(dim=1).eq(topk).all() + assert moe._last_tokens_per_expert.sum().item() == tokens.shape[0] * topk + + fc1, _, fc2 = moe.experts + ref = torch.zeros_like(tokens) + for e in range(num_experts): + w1 = deinterleave_glu_tensor(getattr(fc1, f"weight{e}"), 32) + w2 = getattr(fc2, f"weight{e}") + gate_part, lin_part = (tokens @ w1.t()).chunk(2, dim=-1) + act = torch.nn.functional.silu(gate_part.float()) * lin_part.float() + ref += (act.to(DTYPE) * probs[:, e : e + 1].to(DTYPE)) @ w2.t() + if shared: + ref += moe.shared_expert(tokens) + torch.testing.assert_close(out.reshape(-1, HIDDEN), ref, rtol=0.05, atol=0.05) + + bias_before = moe.expert_bias.clone() + moe.update_expert_bias() + assert torch.isfinite(moe.expert_bias).all() + if topk < num_experts: + assert not torch.equal(bias_before, moe.expert_bias) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index c9b620fa1e..60d17e39af 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -35,6 +35,7 @@ is_bf16_available, ) from transformer_engine.common import recipe +from transformer_engine.pytorch.models import DeepSeekV3Layer from transformer_engine.pytorch.cpp_extensions import general_gemm from transformer_engine.pytorch.tensor.utils import replace_raw_data from transformer_engine.pytorch.module import is_module_grouped_tensor_path_supported @@ -736,6 +737,39 @@ def test_sanity_layernorm_mlp( _test_sanity_common(block, dtype, config, fp8_recipe, skip_wgrad, skip_dgrad, microbatching) +@pytest.mark.parametrize("dtype", param_types) +@pytest.mark.parametrize("fp8_recipe", fp8_recipes, ids=recipe_id) +@pytest.mark.parametrize("moe", all_boolean, ids=["dense", "moe"]) +def test_sanity_deepseek_v3_layer(dtype, fp8_recipe, moe): + config = model_configs["small"] + + if fp8_recipe is not None: + if not is_fp8_supported(config): + pytest.skip("Model config does not support FP8") + if fp8_recipe.nvfp4() and dtype == torch.float16: + pytest.skip("FP16 output for NVFP4 not supported") + + mlp_kwargs = ( + dict(num_experts=4, topk=2, moe_ffn_hidden_size=32, shared_expert_ffn_hidden_size=32) + if moe + else dict(ffn_hidden_size=4 * config.hidden_size) + ) + block = DeepSeekV3Layer( + config.hidden_size, + config.num_heads, + q_lora_rank=16, + kv_lora_rank=16, + qk_nope_head_dim=16, + qk_rope_head_dim=16, + v_head_dim=16, + params_dtype=dtype, + device="cuda", + **mlp_kwargs, + ) + + _test_sanity_e2e(block, dtype, config, fp8_recipe, skip_wgrad=False) + + @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("fp8_recipe", fp8_recipes, ids=recipe_id) @pytest.mark.parametrize("model", ["small"]) diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 576fb57c5e..5940ce7b3d 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -35,6 +35,7 @@ from transformer_engine.pytorch.attention import InferenceParams from transformer_engine.pytorch.attention import RotaryPositionEmbedding from transformer_engine.pytorch.transformer import TransformerLayer +from transformer_engine.pytorch import models from transformer_engine.pytorch.permutation import ( moe_permute, moe_permute_with_probs, diff --git a/transformer_engine/pytorch/models/__init__.py b/transformer_engine/pytorch/models/__init__.py new file mode 100644 index 0000000000..bee5474c81 --- /dev/null +++ b/transformer_engine/pytorch/models/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Model-specific transformer layers composed from Transformer Engine modules.""" + +from transformer_engine.pytorch.models.deepseek_v3 import ( + DeepSeekV3Layer, + DeepSeekV3MoE, + MultiLatentAttention, +) + +__all__ = ["DeepSeekV3Layer", "DeepSeekV3MoE", "MultiLatentAttention"] diff --git a/transformer_engine/pytorch/models/deepseek_v3/__init__.py b/transformer_engine/pytorch/models/deepseek_v3/__init__.py new file mode 100644 index 0000000000..a7cbb50ae2 --- /dev/null +++ b/transformer_engine/pytorch/models/deepseek_v3/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""DeepSeekV3 transformer layer built from Transformer Engine MoE building blocks.""" + +from transformer_engine.pytorch.models.deepseek_v3.multi_latent_attention import ( + MultiLatentAttention, +) +from transformer_engine.pytorch.models.deepseek_v3.moe import DeepSeekV3MoE +from transformer_engine.pytorch.models.deepseek_v3.transformer_layer import DeepSeekV3Layer + +__all__ = ["DeepSeekV3Layer", "DeepSeekV3MoE", "MultiLatentAttention"] diff --git a/tests/pytorch/attention/mla_rope_utils.py b/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py similarity index 63% rename from tests/pytorch/attention/mla_rope_utils.py rename to transformer_engine/pytorch/models/deepseek_v3/mla_rope.py index 90eebfc66a..0aea284afd 100644 --- a/tests/pytorch/attention/mla_rope_utils.py +++ b/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py @@ -2,16 +2,18 @@ # # See LICENSE for license information. -"""MLA RoPE for DSv3 671B - Triton forward and backward kernels. +"""Fused MLA RoPE kernels (DeepSeekV3-style decoupled RoPE/NoPE). -Source: Megatron-LM megatron/core/fusions/fused_mla_yarn_rope_apply.py -Falls back to pure PyTorch when Triton is unavailable. +The query kernel rotates the trailing ``head_dim_rope`` slice in place; the KV +kernel builds the key (nope | broadcast-rotated shared rope head) and value +tensors in a single pass. Falls back to pure PyTorch when Triton is unavailable +or for the ``bshd`` layout. -Note: DSv3 uses YaRN-scaled RoPE for long-context extrapolation. This test -intentionally uses plain RoPE (base=10000) because it only validates MXFP8 -attention path wiring, tensor shapes, forward/backward flow, and relative BF16 -vs MXFP8 behavior. Both reference and MXFP8 paths use the same RoPE tables. -""" +The rope slice is read interleaved (checkpoint layout) and written in NeoX +half-split layout.""" + +import math +from typing import Optional, Tuple import torch @@ -23,25 +25,78 @@ except ImportError: HAVE_TRITON = False -HEAD_DIM_ROPE = 64 -HEAD_DIM_NOPE = 128 -HEAD_DIM_V = 128 -ROTARY_BASE = 10000 +__all__ = [ + "build_rope_tables", + "apply_mla_rope_q", + "apply_mla_rope_kv", + "yarn_mscale", + "yarn_concentration_factor", +] + + +def _yarn_correction_dim(num_rotations, dim, base, max_pos): + return (dim * math.log(max_pos / (num_rotations * 2 * math.pi))) / (2 * math.log(base)) + + +def _yarn_correction_range(beta_fast, beta_slow, dim, base, max_pos, round_to_int=True): + low = _yarn_correction_dim(beta_fast, dim, base, max_pos) + high = _yarn_correction_dim(beta_slow, dim, base, max_pos) + if round_to_int: + low, high = math.floor(low), math.ceil(high) + return max(low, 0), min(high, dim - 1) + + +def _yarn_linear_ramp(low, high, dim, device): + if low == high: + high += 0.001 + ramp = (torch.arange(dim, dtype=torch.float32, device=device) - low) / (high - low) + return torch.clamp(ramp, 0, 1) + + +def yarn_mscale(scale: float, mscale: float = 1.0) -> float: + """YaRN attention temperature factor ``0.1 * mscale * ln(scale) + 1`` (1 for scale <= 1).""" + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +def yarn_concentration_factor(scaling_factor: float, mscale: float, mscale_all_dim: float) -> float: + """Factor multiplied into cos/sin tables.""" + return yarn_mscale(scaling_factor, mscale) / yarn_mscale(scaling_factor, mscale_all_dim) def build_rope_tables( seq_len: int, - emb_dim: int = HEAD_DIM_ROPE, - base: int = ROTARY_BASE, - device: torch.device = None, -) -> tuple[torch.Tensor, torch.Tensor]: - inv_freq = 1.0 / ( - base ** (torch.arange(0, emb_dim, 2, dtype=torch.float32, device=device) / emb_dim) - ) + emb_dim: int, + base: float = 10000.0, + device: Optional[torch.device] = None, + scaling_factor: Optional[float] = None, + original_max_position_embeddings: int = 4096, + beta_fast: float = 32.0, + beta_slow: float = 1.0, + mscale: float = 1.0, + mscale_all_dim: float = 0.0, +) -> Tuple[torch.Tensor, torch.Tensor]: + """cos/sin tables of shape ``[seq_len, emb_dim]`` (fp32, NeoX duplicated halves). + + With ``scaling_factor`` set, frequencies follow YaRN (NTK-by-parts ramp between + ``beta_fast``/``beta_slow`` rotations over ``original_max_position_embeddings``) and the + tables are scaled by the YaRN concentration factor. + """ + exponent = torch.arange(0, emb_dim, 2, dtype=torch.float32, device=device) / emb_dim + inv_freq = 1.0 / (base**exponent) + factor = 1.0 + if scaling_factor is not None: + low, high = _yarn_correction_range( + beta_fast, beta_slow, emb_dim, base, original_max_position_embeddings + ) + extra_mask = 1.0 - _yarn_linear_ramp(low, high, emb_dim // 2, device) + inv_freq = (inv_freq / scaling_factor) * (1 - extra_mask) + inv_freq * extra_mask + factor = yarn_concentration_factor(scaling_factor, mscale, mscale_all_dim) t = torch.arange(seq_len, device=device, dtype=torch.float32) freqs = torch.outer(t, inv_freq) freqs = torch.cat([freqs, freqs], dim=-1) - return torch.cos(freqs).contiguous(), torch.sin(freqs).contiguous() + return (torch.cos(freqs) * factor).contiguous(), (torch.sin(freqs) * factor).contiguous() if HAVE_TRITON: @@ -69,20 +124,9 @@ def _get_thd_token_idx(cu_seqlens, pid_m, seq_num, cp_rank, cp_size): ) * this_seq_len // 2 return token_idx - @triton.autotune( - configs=[ - triton.Config({"BLOCK_H": 1}), - triton.Config({"BLOCK_H": 2}), - triton.Config({"BLOCK_H": 4}), - triton.Config({"BLOCK_H": 8}), - triton.Config({"BLOCK_H": 16}), - triton.Config({"BLOCK_H": 32}), - triton.Config({"BLOCK_H": 64}), - triton.Config({"BLOCK_H": 128}), - ], - key=["emb_dim", "head_num"], - restore_value=["Q"], - ) + _AUTOTUNE_CONFIGS = [triton.Config({"BLOCK_H": h}) for h in (1, 2, 4, 8, 16, 32, 64, 128)] + + @triton.autotune(configs=_AUTOTUNE_CONFIGS, key=["emb_dim", "head_num"], restore_value=["Q"]) @triton.jit def rotary_fwd_q_kernel( Q, @@ -100,6 +144,7 @@ def rotary_fwd_q_kernel( cp_size, BLOCK_H: tl.constexpr, ): + """In-place RoPE fwd on the trailing rope slice of q.""" pid_m = tl.program_id(axis=0) pid_head = tl.program_id(axis=1) if cu_seqlens_q is None: @@ -129,20 +174,7 @@ def rotary_fwd_q_kernel( tl.store(Q + x_left_off, x_left, mask=mask) tl.store(Q + x_right_off, x_right, mask=mask) - @triton.autotune( - configs=[ - triton.Config({"BLOCK_H": 1}), - triton.Config({"BLOCK_H": 2}), - triton.Config({"BLOCK_H": 4}), - triton.Config({"BLOCK_H": 8}), - triton.Config({"BLOCK_H": 16}), - triton.Config({"BLOCK_H": 32}), - triton.Config({"BLOCK_H": 64}), - triton.Config({"BLOCK_H": 128}), - ], - key=["emb_dim", "head_num"], - restore_value=["DO"], - ) + @triton.autotune(configs=_AUTOTUNE_CONFIGS, key=["emb_dim", "head_num"], restore_value=["DO"]) @triton.jit def rotary_bwd_q_kernel( DO, @@ -160,6 +192,7 @@ def rotary_bwd_q_kernel( cp_size, BLOCK_H: tl.constexpr, ): + """In-place RoPE bwd on the trailing rope slice of dq.""" pid_m = tl.program_id(axis=0) pid_head = tl.program_id(axis=1) if cu_seqlens_q is None: @@ -189,19 +222,7 @@ def rotary_bwd_q_kernel( tl.store(DO + x_1_off, x_1, mask=mask) tl.store(DO + x_2_off, x_2, mask=mask) - @triton.autotune( - configs=[ - triton.Config({"BLOCK_H": 1}), - triton.Config({"BLOCK_H": 2}), - triton.Config({"BLOCK_H": 4}), - triton.Config({"BLOCK_H": 8}), - triton.Config({"BLOCK_H": 16}), - triton.Config({"BLOCK_H": 32}), - triton.Config({"BLOCK_H": 64}), - triton.Config({"BLOCK_H": 128}), - ], - key=["emb_dim", "k_dim", "v_dim", "head_num"], - ) + @triton.autotune(configs=_AUTOTUNE_CONFIGS, key=["emb_dim", "k_dim", "v_dim", "head_num"]) @triton.jit def rotary_fwd_kv_kernel( KV, @@ -228,6 +249,7 @@ def rotary_fwd_kv_kernel( cp_size, BLOCK_H: tl.constexpr, ): + """Fwd: build (key, value) from kv and the shared rotated rope head.""" pid_m = tl.program_id(axis=0) pid_head = tl.program_id(axis=1) if cu_seqlens_kv is None: @@ -268,19 +290,7 @@ def rotary_fwd_kv_kernel( tl.store(K_ptr + x_left_off, x_left, mask=mask) tl.store(K_ptr + x_right_off, x_right, mask=mask) - @triton.autotune( - configs=[ - triton.Config({"BLOCK_H": 1}), - triton.Config({"BLOCK_H": 2}), - triton.Config({"BLOCK_H": 4}), - triton.Config({"BLOCK_H": 8}), - triton.Config({"BLOCK_H": 16}), - triton.Config({"BLOCK_H": 32}), - triton.Config({"BLOCK_H": 64}), - triton.Config({"BLOCK_H": 128}), - ], - key=["emb_dim", "k_dim", "v_dim", "head_num"], - ) + @triton.autotune(configs=_AUTOTUNE_CONFIGS, key=["emb_dim", "k_dim", "v_dim", "head_num"]) @triton.jit def rotary_bwd_kv_kernel( dK, @@ -307,6 +317,7 @@ def rotary_bwd_kv_kernel( cp_size, BLOCK_H: tl.constexpr, ): + """Bwd: scatter (dk, dv) into dkv and reduce rope-slice grads into demb.""" pid_m = tl.program_id(axis=0) pid_head = tl.program_id(axis=1) if cu_seqlens_kv is None: @@ -357,19 +368,23 @@ def rotary_bwd_kv_kernel( tl.store(dEMB_ptr + tl.arange(0, emb_dim // 2) * 2, x_1) tl.store(dEMB_ptr + tl.arange(0, emb_dim // 2) * 2 + 1, x_2) - def _flattened_token_stride(tensor: torch.Tensor) -> int: - if tensor.dim() == 4: - return tensor.stride(1) - return tensor.stride(0) + def _token_stride(tensor: torch.Tensor) -> int: + return tensor.stride(1) if tensor.dim() == 4 else tensor.stride(0) class _MLARoPEQTriton(torch.autograd.Function): + """In-place RoPE on the trailing rope slice of q [s, b, h, nope+rope].""" + @staticmethod def forward(ctx, q, cos, sin, head_dim_nope, head_dim_rope): + """Rotate the rope slice of q in place.""" + if not q.is_contiguous(): + q = q.contiguous() s, b, nheads, _ = q.shape - total = s * b - grid_q = lambda META: (total, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_fwd_q_kernel[grid_q]( + def grid(meta): + return (s * b, triton.cdiv(nheads, meta["BLOCK_H"])) + + rotary_fwd_q_kernel[grid]( q, cos, sin, @@ -379,38 +394,38 @@ def forward(ctx, q, cos, sin, head_dim_nope, head_dim_rope): b, None, None, - _flattened_token_stride(q), + _token_stride(q), q.stride(2), 0, 1, ) - ctx.save_for_backward(cos, sin) - ctx.head_dim_nope = head_dim_nope - ctx.head_dim_rope = head_dim_rope - ctx.nheads = nheads - ctx.s = s - ctx.b = b + ctx.dims = (s, b, nheads, head_dim_nope, head_dim_rope) return q @staticmethod def backward(ctx, dq): + """Counter-rotate the rope slice of dq (in place on the copy).""" cos, sin = ctx.saved_tensors - s, b, nheads = ctx.s, ctx.b, ctx.nheads - total = s * b + # attention backward may hand over a strided grad; the kernel + # assumes a contiguous [s, b, h, d] layout + dq = dq.contiguous() + s, b, nheads, head_dim_nope, head_dim_rope = ctx.dims + + def grid(meta): + return (s * b, triton.cdiv(nheads, meta["BLOCK_H"])) - grid_q = lambda META: (total, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_bwd_q_kernel[grid_q]( + rotary_bwd_q_kernel[grid]( dq, cos, sin, - ctx.head_dim_nope, - ctx.head_dim_rope, + head_dim_nope, + head_dim_rope, nheads, b, None, None, - _flattened_token_stride(dq), + _token_stride(dq), dq.stride(2), 0, 1, @@ -418,15 +433,21 @@ def backward(ctx, dq): return dq, None, None, None, None class _MLARoPEKVTriton(torch.autograd.Function): + """kv [s, b, h, nope+v] + shared rope head [s, b, 1, rope] -> (k, v).""" + @staticmethod def forward(ctx, kv, k_pos_emb, cos, sin, head_dim_nope, head_dim_rope, head_dim_v): + """Build (k, v) from kv and the shared rope head.""" + if not kv.is_contiguous(): + kv = kv.contiguous() s, b, nheads, _ = kv.shape - total = s * b - o_key = kv.new_empty(s, b, nheads, head_dim_nope + head_dim_rope) o_value = kv.new_empty(s, b, nheads, head_dim_v) - grid_kv = lambda META: (total, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_fwd_kv_kernel[grid_kv]( + + def grid(meta): + return (s * b, triton.cdiv(nheads, meta["BLOCK_H"])) + + rotary_fwd_kv_kernel[grid]( kv, k_pos_emb, o_key, @@ -440,37 +461,34 @@ def forward(ctx, kv, k_pos_emb, cos, sin, head_dim_nope, head_dim_rope, head_dim b, None, None, - _flattened_token_stride(kv), + _token_stride(kv), kv.stride(2), - _flattened_token_stride(k_pos_emb), - _flattened_token_stride(o_key), + _token_stride(k_pos_emb), + _token_stride(o_key), o_key.stride(2), - _flattened_token_stride(o_value), + _token_stride(o_value), o_value.stride(2), 0, 1, ) - ctx.save_for_backward(cos, sin) - ctx.head_dim_nope = head_dim_nope - ctx.head_dim_rope = head_dim_rope - ctx.head_dim_v = head_dim_v - ctx.nheads = nheads - ctx.s = s - ctx.b = b + ctx.dims = (s, b, nheads, head_dim_nope, head_dim_rope, head_dim_v) return o_key, o_value @staticmethod def backward(ctx, dk_out, dv_out): + """Gradients for (kv, k_pos_emb) from (dk, dv).""" cos, sin = ctx.saved_tensors - s, b, nheads = ctx.s, ctx.b, ctx.nheads - ndp, ndr, ndv = ctx.head_dim_nope, ctx.head_dim_rope, ctx.head_dim_v - total = s * b - + s, b, nheads, ndp, ndr, ndv = ctx.dims + dk_out = dk_out.contiguous() + dv_out = dv_out.contiguous() d_kv = dk_out.new_empty(s, b, nheads, ndp + ndv) d_emb = dk_out.new_empty(s, b, 1, ndr) - grid_kv = lambda META: (total, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_bwd_kv_kernel[grid_kv]( + + def grid(meta): + return (s * b, triton.cdiv(nheads, meta["BLOCK_H"])) + + rotary_bwd_kv_kernel[grid]( dk_out, dv_out, d_kv, @@ -484,185 +502,66 @@ def backward(ctx, dk_out, dv_out): b, None, None, - _flattened_token_stride(dk_out), + _token_stride(dk_out), dk_out.stride(2), - _flattened_token_stride(dv_out), + _token_stride(dv_out), dv_out.stride(2), - _flattened_token_stride(d_kv), + _token_stride(d_kv), d_kv.stride(2), - _flattened_token_stride(d_emb), + _token_stride(d_emb), 0, 1, ) return d_kv, d_emb, None, None, None, None, None -def _apply_mla_rope_q_with_tables( +def _rotate_interleaved_to_neox(x, cos_table, sin_table, seq_dim): + shape = [1, 1, 1, cos_table.shape[-1]] + shape[seq_dim] = cos_table.shape[0] + cos_ = cos_table.view(shape).to(x.dtype) + sin_ = sin_table.view(shape).to(x.dtype) + half = x.shape[-1] // 2 + x_1 = x[..., 0::2] + x_2 = x[..., 1::2] + x_left = x_1 * cos_[..., :half] - x_2 * sin_[..., :half] + x_right = x_2 * cos_[..., half:] + x_1 * sin_[..., half:] + return torch.cat((x_left, x_right), dim=-1) + + +def apply_mla_rope_q( q: torch.Tensor, cos_table: torch.Tensor, sin_table: torch.Tensor, - head_dim_nope: int = HEAD_DIM_NOPE, - head_dim_rope: int = HEAD_DIM_ROPE, + head_dim_nope: int, + head_dim_rope: int, + tensor_format: str = "sbhd", ) -> torch.Tensor: - if HAVE_TRITON: - return _MLARoPEQTriton.apply( - q, - cos_table, - sin_table, - head_dim_nope, - head_dim_rope, - ) - return _apply_pytorch_q(q, cos_table, sin_table, head_dim_nope, head_dim_rope) + """RoPE on the trailing ``head_dim_rope`` slice of q; in place on the Triton path.""" + if HAVE_TRITON and tensor_format == "sbhd": + return _MLARoPEQTriton.apply(q, cos_table, sin_table, head_dim_nope, head_dim_rope) + seq_dim = 0 if tensor_format == "sbhd" else 1 + q_rope = _rotate_interleaved_to_neox(q[..., head_dim_nope:], cos_table, sin_table, seq_dim) + return torch.cat((q[..., :head_dim_nope], q_rope), dim=-1) -def _apply_mla_rope_kv_with_tables( +def apply_mla_rope_kv( kv: torch.Tensor, k_pos_emb: torch.Tensor, cos_table: torch.Tensor, sin_table: torch.Tensor, - head_dim_nope: int = HEAD_DIM_NOPE, - head_dim_rope: int = HEAD_DIM_ROPE, - head_dim_v: int = HEAD_DIM_V, -) -> tuple[torch.Tensor, torch.Tensor]: - if HAVE_TRITON: + head_dim_nope: int, + head_dim_rope: int, + head_dim_v: int, + tensor_format: str = "sbhd", +) -> Tuple[torch.Tensor, torch.Tensor]: + """Build (k, v) from kv ``[.., h, nope+v]`` and the shared rope head ``[.., 1, rope]``.""" + if HAVE_TRITON and tensor_format == "sbhd": return _MLARoPEKVTriton.apply( - kv, - k_pos_emb, - cos_table, - sin_table, - head_dim_nope, - head_dim_rope, - head_dim_v, + kv, k_pos_emb, cos_table, sin_table, head_dim_nope, head_dim_rope, head_dim_v ) - return _apply_pytorch_kv( - kv, - k_pos_emb, - cos_table, - sin_table, - head_dim_nope, - head_dim_rope, - head_dim_v, - ) - - -def apply_mla_rope_q( - q: torch.Tensor, - head_dim_nope: int = HEAD_DIM_NOPE, - head_dim_rope: int = HEAD_DIM_ROPE, - base: int = ROTARY_BASE, - cos_table: torch.Tensor | None = None, - sin_table: torch.Tensor | None = None, -) -> torch.Tensor: - if cos_table is None or sin_table is None: - s = q.shape[0] - cos_table, sin_table = build_rope_tables( - s, - emb_dim=head_dim_rope, - base=base, - device=q.device, - ) - return _apply_mla_rope_q_with_tables( - q, - cos_table, - sin_table, - head_dim_nope, - head_dim_rope, - ) - - -def apply_mla_rope_kv( - kv: torch.Tensor, - k_pos_emb: torch.Tensor, - head_dim_nope: int = HEAD_DIM_NOPE, - head_dim_rope: int = HEAD_DIM_ROPE, - head_dim_v: int = HEAD_DIM_V, - base: int = ROTARY_BASE, - cos_table: torch.Tensor | None = None, - sin_table: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - if cos_table is None or sin_table is None: - s = kv.shape[0] - cos_table, sin_table = build_rope_tables( - s, - emb_dim=head_dim_rope, - base=base, - device=kv.device, - ) - return _apply_mla_rope_kv_with_tables( - kv, - k_pos_emb, - cos_table, - sin_table, - head_dim_nope, - head_dim_rope, - head_dim_v, - ) - - -def apply_mla_rope( - q: torch.Tensor, - kv: torch.Tensor, - k_pos_emb: torch.Tensor, - head_dim_nope: int = HEAD_DIM_NOPE, - head_dim_rope: int = HEAD_DIM_ROPE, - head_dim_v: int = HEAD_DIM_V, - base: int = ROTARY_BASE, - cos_table: torch.Tensor | None = None, - sin_table: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - if cos_table is None or sin_table is None: - s = q.shape[0] - cos_table, sin_table = build_rope_tables( - s, - emb_dim=head_dim_rope, - base=base, - device=q.device, - ) - q = _apply_mla_rope_q_with_tables(q, cos_table, sin_table, head_dim_nope, head_dim_rope) - k, v = _apply_mla_rope_kv_with_tables( - kv, - k_pos_emb, - cos_table, - sin_table, - head_dim_nope, - head_dim_rope, - head_dim_v, - ) - return q, k, v - - -def _rotate_interleaved_to_neox( - x: torch.Tensor, cos_table: torch.Tensor, sin_table: torch.Tensor -) -> torch.Tensor: - cos_ = cos_table[:, None, None, :].to(x.dtype) - sin_ = sin_table[:, None, None, :].to(x.dtype) - half_dim = x.shape[-1] // 2 - x_1 = x[..., 0::2] - x_2 = x[..., 1::2] - x_left = x_1 * cos_[..., :half_dim] - x_2 * sin_[..., :half_dim] - x_right = x_2 * cos_[..., half_dim:] + x_1 * sin_[..., half_dim:] - return torch.cat((x_left, x_right), dim=-1) - - -def _apply_pytorch_q(q, cos_table, sin_table, head_dim_nope, head_dim_rope): - q_nope = q[..., :head_dim_nope] - q_rope = q[..., head_dim_nope : head_dim_nope + head_dim_rope] - q_rope = _rotate_interleaved_to_neox(q_rope, cos_table, sin_table) - return torch.cat((q_nope, q_rope), dim=-1) - - -def _apply_pytorch_kv( - kv, - k_pos_emb, - cos_table, - sin_table, - head_dim_nope, - head_dim_rope, - head_dim_v, -): + seq_dim = 0 if tensor_format == "sbhd" else 1 k_nope = kv[..., :head_dim_nope] v = kv[..., head_dim_nope : head_dim_nope + head_dim_v] - k_rope = _rotate_interleaved_to_neox(k_pos_emb, cos_table, sin_table).expand( - -1, -1, kv.shape[2], -1 - ) - return torch.cat((k_nope, k_rope), dim=-1), v + k_rope = _rotate_interleaved_to_neox(k_pos_emb, cos_table, sin_table, seq_dim) + k_rope = k_rope.expand(*k_nope.shape[:-1], -1) + return torch.cat((k_nope, k_rope), dim=-1), v.contiguous() diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py new file mode 100644 index 0000000000..42f5048e6e --- /dev/null +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -0,0 +1,281 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""DeepSeekV3 MoE block: sigmoid router with aux-loss-free bias, shared + +routed experts.""" + +from typing import Optional, Union + +import torch + +import transformer_engine.pytorch.ops as te_ops +from transformer_engine.pytorch.router import fused_topk_with_score_function +from transformer_engine.pytorch.permutation import ( + moe_permute_and_pad_with_probs, + moe_permute_with_probs, + moe_unpermute, +) +from transformer_engine.pytorch.quantization import ( + FP8GlobalStateManager, + get_align_size_for_quantization, +) + +__all__ = ["DeepSeekV3MoE"] + + +_EP_ALIGNMENT = 128 + + +def _make_swiglu_mlp(hidden_size, ffn_hidden_size, dtype, device, num_experts=None): + """Dense SwiGLU MLP, or a grouped one (probs applied inside the activation) per expert. + + The grouped variant fuses into a single CuTe grouped MLP on supported hardware. + """ + common = {"bias": False, "dtype": dtype, "device": device} + if num_experts is None: + return te_ops.Sequential( + te_ops.Linear(hidden_size, 2 * ffn_hidden_size, **common), + te_ops.SwiGLU(), + te_ops.Linear(ffn_hidden_size, hidden_size, **common), + ) + return te_ops.Sequential( + te_ops.GroupedLinear(num_experts, hidden_size, 2 * ffn_hidden_size, **common), + te_ops.ScaledSwiGLU(glu_interleave_size=32), + te_ops.GroupedLinear(num_experts, ffn_hidden_size, hidden_size, **common), + ) + + +class DeepSeekV3MoE(torch.nn.Module): + """ + DeepSeekV3 Mixture-of-Experts block. + + Each token is scored by a sigmoid router with a non-trainable expert bias + updated by ``update_expert_bias()`` (aux-loss-free load balancing) and, + optionally, group-limited routing: experts are split into ``num_groups`` + groups, the top ``group_topk`` groups are selected by their summed scores, + and the final ``topk`` experts are chosen only from those groups. Selected + tokens run through the routed experts, a SwiGLU MLP shared across experts + as a grouped GEMM, with the routing probability applied inside the MLP. An + optional shared expert (dense SwiGLU MLP) is added to every token. On + hardware that supports it the expert MLP runs as a single fused + grouped-GEMM kernel. + + Without ``ep_group`` all experts live on the local device. With + ``ep_group`` the experts are split across the group and tokens are + exchanged over NCCL; this requires ``ep_bootstrap`` to be called once per + process before constructing the module, and bfloat16 inputs. + + Parameters + ---------- + hidden_size : int + size of each input sample. + moe_ffn_hidden_size : int + ffn size of each routed expert. + num_experts : int + total number of routed experts. + topk : int, default = 8 + number of experts per token. + num_groups : int, optional + number of expert groups for node-limited routing. + group_topk : int, optional + number of groups each token is limited to. + routed_scaling_factor : float, default = 2.5 + scaling applied to the routing probabilities. + shared_expert_ffn_hidden_size : int, optional + ffn size of the shared expert; ``None`` + disables the shared expert. + expert_bias_update_rate : float, default = 1e-3 + step size of the aux-loss-free bias update + (see :meth:`update_expert_bias`). + params_dtype : torch.dtype, optional + dtype of module parameters. + ep_group : ProcessGroup, optional + expert-parallel process group; enables the NCCL EP path. + ep_max_tokens_per_rank : int, optional + max local tokens per forward (required with EP). + """ + + def __init__( + self, + hidden_size: int, + moe_ffn_hidden_size: int, + num_experts: int, + topk: int = 8, + num_groups: Optional[int] = None, + group_topk: Optional[int] = None, + routed_scaling_factor: float = 2.5, + shared_expert_ffn_hidden_size: Optional[int] = None, + expert_bias_update_rate: float = 1e-3, + params_dtype: Optional[torch.dtype] = None, + device: Union[torch.device, str] = "cuda", + ep_group: Optional[torch.distributed.ProcessGroup] = None, + ep_max_tokens_per_rank: Optional[int] = None, + ) -> None: + super().__init__() + + dtype = params_dtype if params_dtype is not None else torch.get_default_dtype() + self.hidden_size = hidden_size + self.num_experts = num_experts + self.topk = topk + self.num_groups = num_groups + self.group_topk = group_topk + self.routed_scaling_factor = routed_scaling_factor + self.expert_bias_update_rate = expert_bias_update_rate + + self.gate = torch.nn.Linear( + hidden_size, num_experts, bias=False, dtype=dtype, device=device + ) + self.register_buffer( + "expert_bias", torch.zeros(num_experts, dtype=torch.float32, device=device) + ) + self._last_tokens_per_expert: Optional[torch.Tensor] = None + + self.ep_group = ep_group + self.ep_size = 1 if ep_group is None else torch.distributed.get_world_size(ep_group) + assert num_experts % self.ep_size == 0 + num_local_experts = num_experts // self.ep_size + + self.experts = _make_swiglu_mlp( + hidden_size, moe_ffn_hidden_size, dtype, device, num_experts=num_local_experts + ) + + self.shared_expert = None + if shared_expert_ffn_hidden_size is not None: + self.shared_expert = _make_swiglu_mlp( + hidden_size, shared_expert_ffn_hidden_size, dtype, device + ) + + self.ep_buffer = None + if ep_group is not None: + from transformer_engine.pytorch.ep import EpBuffer + + assert ep_max_tokens_per_rank is not None, "EP requires ep_max_tokens_per_rank." + # Worst case plus per-expert alignment padding, rounded up to + # the multiple of 128 required by the fused grouped MLP. + cap = self.ep_size * ep_max_tokens_per_rank * topk + cap += num_local_experts * _EP_ALIGNMENT + cap = -(-cap // _EP_ALIGNMENT) * _EP_ALIGNMENT + self.ep_buffer = EpBuffer( + top_k=topk, + max_tokens_per_rank=ep_max_tokens_per_rank, + hidden_dim=hidden_size, + num_local_experts=num_local_experts, + recv_capacity_per_rank=cap, + alignment=_EP_ALIGNMENT, + device=device, + ) + + def _route(self, logits: torch.Tensor, topk_indices: Optional[torch.Tensor] = None): + return fused_topk_with_score_function( + logits=logits, + topk=self.topk, + use_pre_softmax=False, + num_groups=self.num_groups, + group_topk=self.group_topk, + scaling_factor=self.routed_scaling_factor, + score_function="sigmoid", + expert_bias=self.expert_bias, + topk_indices=topk_indices, + ) + + def _forward_local(self, tokens: torch.Tensor) -> torch.Tensor: + probs, routing_map = self._route(self.gate(tokens).float()) + tokens_per_expert = routing_map.sum(dim=0) + self._last_tokens_per_expert = tokens_per_expert.detach() + + # Quantized grouped GEMMs need every expert's row count aligned. + align = 1 + if FP8GlobalStateManager.is_fp8_enabled(): + align = get_align_size_for_quantization(FP8GlobalStateManager.get_fp8_recipe()) + if align > 1: + permuted, permuted_probs, row_id_map, pad_offsets, tokens_per_expert = ( + moe_permute_and_pad_with_probs(tokens, probs, routing_map, tokens_per_expert, align) + ) + else: + permuted, permuted_probs, row_id_map = moe_permute_with_probs( + tokens, probs, routing_map, num_out_tokens=tokens.shape[0] * self.topk + ) + pad_offsets = None + + # The fused grouped MLP requires the total row count to be a multiple + # of 128; rows beyond sum(tokens_per_expert) fall outside every group. + num_rows = permuted.shape[0] + pad = (-num_rows) % 128 + if pad: + permuted = torch.nn.functional.pad(permuted, (0, 0, 0, pad)) + permuted_probs = torch.nn.functional.pad(permuted_probs, (0, pad)) + + out = self.experts( + permuted, tokens_per_expert, permuted_probs.to(tokens.dtype), tokens_per_expert + ) + return moe_unpermute( + out[:num_rows], row_id_map, restore_shape=tokens.shape, pad_offsets=pad_offsets + ) + + def _forward_ep(self, tokens: torch.Tensor) -> torch.Tensor: + from transformer_engine.pytorch.ep import ep_dispatch, ep_combine + + assert tokens.dtype == torch.bfloat16, "The EP path requires bfloat16 inputs." + topk_idx = torch.empty( + (tokens.shape[0], self.topk), dtype=torch.int64, device=tokens.device + ) + probs, topk_idx = self._route(self.gate(tokens).float(), topk_indices=topk_idx) + flat_idx = topk_idx.flatten() + self._last_tokens_per_expert = torch.zeros( + self.num_experts, dtype=torch.long, device=tokens.device + ).scatter_add_(0, flat_idx, torch.ones_like(flat_idx)) + topk_weights = probs.gather(1, topk_idx) + + # Zero-filled recv/grad buffers: per-expert alignment padding lands + # inside the grouped-GEMM m_splits, so uninitialized rows would poison + # the expert wgrads. + cap = self.ep_buffer.recv_capacity_per_rank + recv_tokens, recv_weights, tokens_per_expert = ep_dispatch( + self.ep_buffer, + tokens, + topk_idx, + topk_weights, + recv_tokens=torch.zeros( + (cap, self.hidden_size), dtype=tokens.dtype, device=tokens.device + ), + recv_topk_weights=torch.zeros((cap,), dtype=torch.float32, device=tokens.device), + ) + expert_out = self.experts( + recv_tokens, tokens_per_expert, recv_weights.to(tokens.dtype), tokens_per_expert + ) + return ep_combine( + self.ep_buffer, + expert_out, + num_local_tokens=tokens.shape[0], + grad_out=torch.zeros_like(expert_out), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """ + Parameters + ---------- + hidden_states : torch.Tensor + input of shape ``[..., hidden_size]``. + """ + tokens = hidden_states.reshape(-1, self.hidden_size) + if self.ep_group is not None: + out = self._forward_ep(tokens) + else: + out = self._forward_local(tokens) + if self.shared_expert is not None: + out = out + self.shared_expert(tokens) + return out.view_as(hidden_states) + + @torch.no_grad() + def update_expert_bias(self) -> None: + """Aux-loss-free bias update from the last forward's routing counts. + + With data/expert parallelism, all-reduce ``_last_tokens_per_expert`` + across ranks before calling (or call on identically-routed ranks). + """ + counts = self._last_tokens_per_expert + if counts is None: + return + err = counts.float().mean() - counts.float() + self.expert_bias += self.expert_bias_update_rate * torch.sign(err) diff --git a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py new file mode 100644 index 0000000000..56b4d3d0d1 --- /dev/null +++ b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py @@ -0,0 +1,257 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Multi-Latent Attention (MLA) block as used in DeepSeekV3.""" + +import math +from typing import Optional, Union + +import torch + +from transformer_engine.pytorch.module import Linear, LayerNormLinear +from transformer_engine.pytorch.attention import DotProductAttention +from transformer_engine.pytorch.models.deepseek_v3.mla_rope import ( + apply_mla_rope_kv, + apply_mla_rope_q, + build_rope_tables, + yarn_mscale, +) + +__all__ = ["MultiLatentAttention"] + + +class MultiLatentAttention(torch.nn.Module): + """ + Multi-Latent Attention as used in DeepSeekV3. + + Queries and key-values are projected through low-rank latents + (``q_lora_rank``, ``kv_lora_rank``); RMSNorm on each latent is fused into + the up-projection (:class:`LayerNormLinear` with RMSNorm). Each query/key + head is split into a ``qk_nope_head_dim`` part and a ``qk_rope_head_dim`` + part; RoPE is applied only to the rope part, and the key rope part comes + from a single shared head broadcast to all heads. Attention runs through + :class:`DotProductAttention` with asymmetric head dims + ``kv_channels=(qk_nope_head_dim + qk_rope_head_dim, v_head_dim)``, which + supports the cuDNN fused attention backend. + + RoPE uses the fused MLA kernels from :mod:`.mla_rope` (in-place on the + query rope slice, single-pass key/value assembly); the rope slice follows + the DeepSeekV3 checkpoint convention (interleaved weights, NeoX output). + + Parameters + ---------- + hidden_size : int + size of each input sample. + num_attention_heads : int + number of attention heads. + q_lora_rank : int, default = 1536 + rank of the query latent. + kv_lora_rank : int, default = 512 + rank of the key-value latent. + qk_nope_head_dim : int, default = 128 + per-head dim of the non-rotary query/key part. + qk_rope_head_dim : int, default = 64 + per-head dim of the rotary query/key part. + v_head_dim : int, default = 128 + per-head dim of the values. + attention_dropout : float, default = 0.0 + dropout probability on attention scores. + attn_mask_type : str, default = "causal" + attention mask type passed to :class:`DotProductAttention`. + layernorm_epsilon : float, default = 1e-6 + epsilon of the latent RMSNorms (matches DeepSeekV3). + rotary_base : float, default = 10000.0 + RoPE base. + rope_scaling_factor : float, optional + YaRN context-extension factor; ``None`` disables YaRN. + original_max_position_embeddings : int, default = 4096 + pre-extension context length (YaRN). + beta_fast : float, default = 32.0 + YaRN high-frequency rotation bound. + beta_slow : float, default = 1.0 + YaRN low-frequency rotation bound. + mscale : float, default = 1.0 + YaRN mscale of the rope part. + mscale_all_dim : float, default = 0.0 + YaRN mscale of all dims; sets the default softmax scale to + ``m**2 / sqrt(qk head dim)`` with ``m = 0.1 * mscale_all_dim * ln(factor) + 1``. + softmax_scale : float, optional + softmax scale; defaults to ``1/sqrt(qk head dim)`` (times the YaRN + ``m**2`` when YaRN is enabled). + qkv_format : str, default = "sbhd" + layout of the input/output tensors. + params_dtype : torch.dtype, optional + dtype of module parameters. + tp_group : ProcessGroup, optional + tensor-parallel process group for the up/output projections. + tp_size : int, default = 1 + tensor-parallel world size. + """ + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + q_lora_rank: int = 1536, + kv_lora_rank: int = 512, + qk_nope_head_dim: int = 128, + qk_rope_head_dim: int = 64, + v_head_dim: int = 128, + attention_dropout: float = 0.0, + attn_mask_type: str = "causal", + layernorm_epsilon: float = 1e-6, + rotary_base: float = 10000.0, + rope_scaling_factor: Optional[float] = None, + original_max_position_embeddings: int = 4096, + beta_fast: float = 32.0, + beta_slow: float = 1.0, + mscale: float = 1.0, + mscale_all_dim: float = 0.0, + softmax_scale: Optional[float] = None, + qkv_format: str = "sbhd", + params_dtype: Optional[torch.dtype] = None, + tp_group: Optional[torch.distributed.ProcessGroup] = None, + tp_size: int = 1, + device: Union[torch.device, str] = "cuda", + ) -> None: + super().__init__() + + assert qkv_format in ("sbhd", "bshd"), "MultiLatentAttention supports sbhd/bshd formats." + assert num_attention_heads % tp_size == 0 + + self.qkv_format = qkv_format + self.num_attention_heads = num_attention_heads + self.num_attention_heads_per_partition = num_attention_heads // tp_size + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.v_head_dim = v_head_dim + self.kv_lora_rank = kv_lora_rank + + common = {"bias": False, "params_dtype": params_dtype, "device": device} + tp = {"tp_group": tp_group, "tp_size": tp_size} + + self.q_down_proj = Linear(hidden_size, q_lora_rank, **common) + self.q_up_proj = LayerNormLinear( + q_lora_rank, + num_attention_heads * self.qk_head_dim, + normalization="RMSNorm", + eps=layernorm_epsilon, + parallel_mode="column" if tp_size > 1 else None, + **tp, + **common, + ) + self.kv_down_proj = Linear(hidden_size, kv_lora_rank + qk_rope_head_dim, **common) + self.kv_up_proj = LayerNormLinear( + kv_lora_rank, + num_attention_heads * (qk_nope_head_dim + v_head_dim), + normalization="RMSNorm", + eps=layernorm_epsilon, + parallel_mode="column" if tp_size > 1 else None, + **tp, + **common, + ) + self.out_proj = Linear( + num_attention_heads * v_head_dim, + hidden_size, + parallel_mode="row" if tp_size > 1 else None, + **tp, + **common, + ) + + self.rotary_base = rotary_base + self._yarn_kwargs = { + "scaling_factor": rope_scaling_factor, + "original_max_position_embeddings": original_max_position_embeddings, + "beta_fast": beta_fast, + "beta_slow": beta_slow, + "mscale": mscale, + "mscale_all_dim": mscale_all_dim, + } + self._rope_tables: Optional[tuple] = None + + if softmax_scale is None and rope_scaling_factor is not None: + m = yarn_mscale(rope_scaling_factor, mscale_all_dim) + softmax_scale = m * m / math.sqrt(self.qk_head_dim) + self.softmax_scale = softmax_scale + + self.core_attention = DotProductAttention( + num_attention_heads, + kv_channels=(self.qk_head_dim, v_head_dim), + attention_dropout=attention_dropout, + qkv_format=qkv_format, + attn_mask_type=attn_mask_type, + softmax_scale=softmax_scale, + tp_group=tp_group, + tp_size=tp_size, + ) + + def _rope_tables_for(self, seq_len: int, device: torch.device): + if self._rope_tables is None or self._rope_tables[0].shape[0] < seq_len: + self._rope_tables = build_rope_tables( + seq_len, + self.qk_rope_head_dim, + base=self.rotary_base, + device=device, + **self._yarn_kwargs, + ) + cos, sin = self._rope_tables + return cos[:seq_len], sin[:seq_len] + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + attn_mask_type: Optional[str] = None, + checkpoint_core_attention: bool = False, + ) -> torch.Tensor: + """ + Parameters + ---------- + hidden_states : torch.Tensor + input of shape ``[sq, b, h]`` (sbhd) or ``[b, sq, h]`` (bshd). + attention_mask : torch.Tensor, optional + boolean mask passed to :class:`DotProductAttention`. + attn_mask_type : str, optional + override of the constructor's mask type. + checkpoint_core_attention : bool, default = False + checkpoint the core attention computation. + """ + seq_dim = 0 if self.qkv_format == "sbhd" else 1 + seq_len = hidden_states.shape[seq_dim] + heads = self.num_attention_heads_per_partition + + q = self.q_up_proj(self.q_down_proj(hidden_states)) + q = q.view(*q.shape[:-1], heads, self.qk_head_dim) + + kv_down = self.kv_down_proj(hidden_states) + kv_latent, k_pos = torch.split(kv_down, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + kv = self.kv_up_proj(kv_latent) + kv = kv.view(*kv.shape[:-1], heads, self.qk_nope_head_dim + self.v_head_dim) + + cos, sin = self._rope_tables_for(seq_len, hidden_states.device) + q = apply_mla_rope_q( + q, cos, sin, self.qk_nope_head_dim, self.qk_rope_head_dim, self.qkv_format + ) + k, v = apply_mla_rope_kv( + kv, + k_pos.unsqueeze(-2), + cos, + sin, + self.qk_nope_head_dim, + self.qk_rope_head_dim, + self.v_head_dim, + self.qkv_format, + ) + + context = self.core_attention( + q, + k, + v, + attention_mask=attention_mask, + qkv_format=self.qkv_format, + attn_mask_type=attn_mask_type, + checkpoint_core_attention=checkpoint_core_attention, + ) + return self.out_proj(context) diff --git a/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py new file mode 100644 index 0000000000..ab11fe6394 --- /dev/null +++ b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py @@ -0,0 +1,179 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""DeepSeekV3 transformer layer.""" + +from typing import Optional, Union + +import torch + +from transformer_engine.pytorch.module import LayerNormMLP, RMSNorm +from transformer_engine.pytorch.models.deepseek_v3.multi_latent_attention import ( + MultiLatentAttention, +) +from transformer_engine.pytorch.models.deepseek_v3.moe import DeepSeekV3MoE + +__all__ = ["DeepSeekV3Layer"] + + +class DeepSeekV3Layer(torch.nn.Module): + """ + A full DeepSeekV3 transformer layer, analogous to + :class:`TransformerLayer`: pre-RMSNorm + :class:`MultiLatentAttention`, + then either a dense SwiGLU MLP (:class:`LayerNormMLP` with RMSNorm, used + for the first dense layers of DeepSeekV3) or :class:`DeepSeekV3MoE`, each + with a residual connection. + + Parameters + ---------- + hidden_size : int + size of each input sample. + num_attention_heads : int + number of attention heads. + ffn_hidden_size : int + ffn size of the dense MLP (used when ``num_experts`` is + ``None``). + num_experts : int, optional + number of routed experts; ``None`` makes this a dense layer. + moe_ffn_hidden_size : int, optional + ffn size of each routed expert (required with MoE). + hidden_dropout : float, default = 0.0 + dropout probability on the residual branches. + **kwargs + kwargs common to the submodules (``q_lora_rank``, ``kv_lora_rank``, + ``qk_nope_head_dim``, ``qk_rope_head_dim``, ``v_head_dim``, + ``attention_dropout``, ``attn_mask_type``, ``qkv_format``, ``topk``, + ``num_groups``, ``group_topk``, ``routed_scaling_factor``, + ``shared_expert_ffn_hidden_size``, EP options, ...), forwarded to + :class:`MultiLatentAttention` and :class:`DeepSeekV3MoE`. + """ + + _MLA_KWARGS = frozenset( + { + "q_lora_rank", + "kv_lora_rank", + "qk_nope_head_dim", + "qk_rope_head_dim", + "v_head_dim", + "attention_dropout", + "attn_mask_type", + "rotary_base", + "rope_scaling_factor", + "original_max_position_embeddings", + "beta_fast", + "beta_slow", + "mscale", + "mscale_all_dim", + "softmax_scale", + "qkv_format", + "tp_group", + "tp_size", + } + ) + _MOE_KWARGS = frozenset( + { + "topk", + "num_groups", + "group_topk", + "routed_scaling_factor", + "shared_expert_ffn_hidden_size", + "expert_bias_update_rate", + "ep_group", + "ep_max_tokens_per_rank", + } + ) + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + ffn_hidden_size: Optional[int] = None, + num_experts: Optional[int] = None, + moe_ffn_hidden_size: Optional[int] = None, + hidden_dropout: float = 0.0, + layernorm_epsilon: float = 1e-5, + params_dtype: Optional[torch.dtype] = None, + device: Union[torch.device, str] = "cuda", + **kwargs, + ) -> None: + super().__init__() + + unknown = set(kwargs) - self._MLA_KWARGS - self._MOE_KWARGS + if unknown: + raise TypeError(f"Unexpected keyword arguments: {sorted(unknown)}") + mla_kwargs = {k: v for k, v in kwargs.items() if k in self._MLA_KWARGS} + moe_kwargs = {k: v for k, v in kwargs.items() if k in self._MOE_KWARGS} + + self.hidden_dropout = hidden_dropout + + self.input_layernorm = RMSNorm( + hidden_size, eps=layernorm_epsilon, device=device, dtype=params_dtype + ) + self.self_attention = MultiLatentAttention( + hidden_size, + num_attention_heads, + params_dtype=params_dtype, + device=device, + **mla_kwargs, + ) + + if num_experts is None: + assert ffn_hidden_size is not None, "Dense layers require ffn_hidden_size." + self.pre_mlp_layernorm = None + self.mlp = LayerNormMLP( + hidden_size, + ffn_hidden_size, + eps=layernorm_epsilon, + normalization="RMSNorm", + activation="swiglu", + bias=False, + params_dtype=params_dtype, + device=device, + ) + else: + assert moe_ffn_hidden_size is not None, "MoE layers require moe_ffn_hidden_size." + self.pre_mlp_layernorm = RMSNorm( + hidden_size, eps=layernorm_epsilon, device=device, dtype=params_dtype + ) + self.mlp = DeepSeekV3MoE( + hidden_size, + moe_ffn_hidden_size, + num_experts, + params_dtype=params_dtype, + device=device, + **moe_kwargs, + ) + + def _residual_add(self, out: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: + out = torch.nn.functional.dropout(out, p=self.hidden_dropout, training=self.training) + return residual + out + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + checkpoint_core_attention: bool = False, + ) -> torch.Tensor: + """ + Parameters + ---------- + hidden_states : torch.Tensor + input of shape ``[sq, b, h]`` (sbhd) or ``[b, sq, h]`` (bshd). + attention_mask : torch.Tensor, optional + boolean attention mask. + checkpoint_core_attention : bool, default = False + checkpoint the core attention computation. + """ + attention_out = self.self_attention( + self.input_layernorm(hidden_states), + attention_mask=attention_mask, + checkpoint_core_attention=checkpoint_core_attention, + ) + hidden_states = self._residual_add(attention_out, hidden_states) + + if self.pre_mlp_layernorm is not None: + mlp_out = self.mlp(self.pre_mlp_layernorm(hidden_states)) + else: + mlp_out = self.mlp(hidden_states) + return self._residual_add(mlp_out, hidden_states)