diff --git a/docs/fern/versions/nightly.yml b/docs/fern/versions/nightly.yml index 088935d639..3c83de1811 100644 --- a/docs/fern/versions/nightly.yml +++ b/docs/fern/versions/nightly.yml @@ -268,6 +268,9 @@ navigation: - page: "MiniMax-M3" path: ../../model-coverage/vlm/minimax/minimax-m3.mdx slug: minimax-m3 + - page: "GLM-5.3-Flash" + path: ../../model-coverage/vlm/thudm/glm-5-3-flash.mdx + slug: glm-5-3-flash - section: "Multimodal" slug: multimodal contents: diff --git a/docs/model-coverage/vlm/thudm/glm-5-3-flash.mdx b/docs/model-coverage/vlm/thudm/glm-5-3-flash.mdx new file mode 100644 index 0000000000..87fa3339bc --- /dev/null +++ b/docs/model-coverage/vlm/thudm/glm-5-3-flash.mdx @@ -0,0 +1,41 @@ +--- +title: "GLM-5.3-Flash" +description: "Fine-tune the GLM-5.3-Flash mixture-of-experts vision-language model with packed context and expert parallelism." +slug: model-coverage/vision-language-models/thudm/glm-5-3-flash +--- + +[GLM-5.3-Flash](https://huggingface.co/zai-org/GLM-5.3-Flash) is a mixture-of-experts vision-language model with a hybrid Kimi Delta Attention and Dynamic Sparse Attention language backbone. + + + +| | | +|---|---| +| **Task** | Image-Text-to-Text | +| **Architecture** | `Glm5NextForConditionalGeneration` | +| **Language Module** | Hybrid KDA / KPool-DSA MoE decoder | +| **Training Precision** | BF16 after FP8 checkpoint dequantization | +| **HF Org** | [zai-org](https://huggingface.co/zai-org) | + + + +## Supported Training Path + +NeMo AutoModel provides a native configuration, image processor, vision tower, language model, and Hugging Face state-dict adapter for GLM-5.3-Flash. The supported full-model initialization path loads the base checkpoint through distributed checkpointing; single-GPU full-checkpoint loading is not supported. + +The current onboarding supports image training. Video inputs, tensor parallelism, and pipeline parallelism are not enabled for this model. The validated recipe uses FSDP2 with expert parallelism, contiguous packed context parallelism, and HybridEP dispatch. + +## Attention Backends + +- KDA layers use Flash Linear Attention kernels. +- Sparse DSA layers support the SDPA reference path. +- On SM90 or later, `backend.attn: cudnn` uses FlashMLA forward with cuDNN sparse-attention backward. This optional path requires compatible FlashMLA and cuDNN Frontend installations. + +## Example Recipe + +- [Full SFT — MedPix, packed 2K, EP72 + CP2](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/vlm_finetune/glm5_next/glm5_3_flash_medpix_packed2k_ep72_cp2_100steps.yaml) + +The recipe is sized for 9 nodes with 8 GPUs per node. It uses local batch size 1 and four gradient-accumulation microsteps to form a global batch of 144 packed samples. + +## Hugging Face Model Card + +- [zai-org/GLM-5.3-Flash](https://huggingface.co/zai-org/GLM-5.3-Flash) diff --git a/examples/vlm_finetune/glm5_next/glm5_3_flash_medpix_packed2k_ep72_cp2_100steps.yaml b/examples/vlm_finetune/glm5_next/glm5_3_flash_medpix_packed2k_ep72_cp2_100steps.yaml new file mode 100644 index 0000000000..087f084da7 --- /dev/null +++ b/examples/vlm_finetune/glm5_next/glm5_3_flash_medpix_packed2k_ep72_cp2_100steps.yaml @@ -0,0 +1,155 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Validated 9-node/72-GPU packed MedPix run for GLM-5.3-Flash. CP2 gives DP36; +# EP72 spans the dp(36) x cp(2) mesh and stores four of 288 routed experts per +# GPU. Four accumulation microsteps form each 144-pack optimizer step. TP and PP +# remain disabled. + +recipe: FinetuneRecipeForVLM +seed: 1234 + +step_scheduler: + global_batch_size: 144 + local_batch_size: 1 + ckpt_every_steps: 1000 + val_every_steps: 1000 + log_remote_every_steps: 1 + gc_every_steps: 5 + num_epochs: 4 + max_steps: 100 + +dist_env: + backend: nccl + timeout_minutes: 120 + +rng: + _target_: nemo_automodel.components.training.rng.StatefulRNG + seed: 1234 + ranked: true + deterministic: true + +model: + _target_: nemo_automodel.NeMoAutoModelForImageTextToText.from_pretrained + pretrained_model_name_or_path: zai-org/GLM-5.3-Flash + torch_dtype: bfloat16 + attn_implementation: sdpa + use_liger_kernel: false + use_sdpa_patching: false + text_config: + output_hidden_states: true + num_nextn_predict_layers: 0 + backend: + _target_: nemo_automodel.components.models.common.BackendConfig + # KDA layers remain on FLA; this selects cuDNN/FlashMLA for sparse MLA layers. + attn: cudnn + linear: torch + rms_norm: torch_fp32 + experts: torch_mm + dispatcher: hybridep + rope_fusion: false + gate_precision: float32 + fake_balanced_gate: false + enable_hf_state_dict_adapter: true + enable_fsdp_optimizations: true + +processor: + _target_: nemo_automodel.components.models.glm5_next.processing.build_glm5_next_processor + pretrained_model_name_or_path: zai-org/GLM-5.3-Flash + trust_remote_code: false + +checkpoint: + enabled: false + checkpoint_dir: checkpoints/glm5_3_flash_medpix_packed2k_ep72_cp2_100steps/ + model_save_format: safetensors + save_consolidated: false + dequantize_base_checkpoint: true + +distributed: + strategy: fsdp2 + tp_size: 1 + cp_size: 2 + pp_size: 1 + ep_size: 72 + sequence_parallel: false + activation_checkpointing: true + # Avoid retaining a full accumulated gradient set across four microsteps. + defer_fsdp_grad_sync: false + moe: + reshard_after_forward: false + wrap_outer_model: true + ignore_router_for_ac: true + +freeze_config: + freeze_embeddings: true + freeze_vision_tower: true + freeze_audio_tower: true + freeze_language_model: false + +loss_fn: + _target_: nemo_automodel.components.loss.linear_ce.FusedLinearCrossEntropy + +dataset: + _target_: nemo_automodel.components.datasets.vlm.datasets.make_medpix_dataset + path_or_dataset: mmoukouba/MedPix-VQA + split: train + +packed_sequence: + pretokenize: true + max_length: 2048 + pack_size: 2048 + collate_max_length: 2048 + packing_ratio: 0.9 + drop_long_samples: true + balance_media_tokens: true + packing_format: neat + attn_implementation: sdpa + +dataloader: + _target_: torchdata.stateful_dataloader.StatefulDataLoader + num_workers: 1 + persistent_workers: true + pin_memory: true + drop_last: true + +validation_dataset: none +validation_dataloader: none + +optimizer: + _target_: torch.optim.AdamW + betas: [0.9, 0.95] + eps: 1e-8 + lr: 5.0e-6 + weight_decay: 0.1 + +lr_scheduler: + lr_decay_style: constant + +clip_grad_norm: + max_norm: 1.0 + +wandb: + enable: false + project: huiyingl_workspace + entity: Nemo-automodel + name: glm5_3_flash_medpix_packed2k_ep72_cp2_100steps + group: glm5_3_flash_medpix_packed2k_cp_parity_100steps + tags: [glm5.3-flash, medpix, vlm, packed2k, ep72, cp2, parity] + dir: logs/glm5_3_flash_medpix_packed2k_ep72_cp2_100steps/wandb + +ci: + recipe_owner: HuiyingLi + nodes: 9 + time: "01:00:00" + max_steps: 5 diff --git a/nemo_automodel/_transformers/registry.py b/nemo_automodel/_transformers/registry.py index db8bcfaa56..d06a50f764 100644 --- a/nemo_automodel/_transformers/registry.py +++ b/nemo_automodel/_transformers/registry.py @@ -86,6 +86,13 @@ "GlmMoeDsaForCausalLM", ("nemo_automodel.components.models.glm_moe_dsa.model", "GlmMoeDsaForCausalLM"), ), + ( + "Glm5NextForConditionalGeneration", + ( + "nemo_automodel.components.models.glm5_next.model", + "Glm5NextForConditionalGeneration", + ), + ), ( "Gemma4ForConditionalGeneration", ("nemo_automodel.components.models.gemma4_moe.model", "Gemma4ForConditionalGeneration"), @@ -326,6 +333,7 @@ "bailing_moe": ("nemo_automodel.components.models.ling_v2.config", "BailingMoeV2Config"), "deepseek_v4": ("nemo_automodel.components.models.deepseek_v4.config", "DeepseekV4Config"), "glm_moe_dsa": ("nemo_automodel.components.models.glm_moe_dsa.config", "GlmMoeDsaConfig"), + "glm5_next": ("nemo_automodel.components.models.glm5_next.config", "Glm5NextConfig"), "hy_v3": ("nemo_automodel.components.models.hy_v3.config", "HYV3Config"), "inkling_mm_model": ("nemo_automodel.components.models.inkling.configuration", "InklingConfig"), "kimi_k2": ("nemo_automodel.components.models.kimi_k2.config", "KimiK2Config"), diff --git a/nemo_automodel/components/models/common/cudnn_sparse_attention.py b/nemo_automodel/components/models/common/cudnn_sparse_attention.py new file mode 100644 index 0000000000..14b33bcc66 --- /dev/null +++ b/nemo_automodel/components/models/common/cudnn_sparse_attention.py @@ -0,0 +1,400 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared FlashMLA-forward/cuDNN-backward sparse latent attention.""" + +from __future__ import annotations + +import math +from typing import Any + +import torch + +from nemo_automodel.shared.import_utils import safe_import_from + +_HAS_CUDNN_DSA, _CUDNN_DSA = safe_import_from( + "cudnn", + "DSA", + msg=( + "cuDNN sparse-attention kernels are unavailable. Install " + "nvidia-cudnn-frontend[cutedsl] to use backend.attn='cudnn'." + ), +) +_HAS_FLASH_MLA, _FLASH_MLA_SPARSE_FWD = safe_import_from( + "flash_mla", + "flash_mla_sparse_fwd", + msg="FlashMLA sparse prefill is unavailable. Install the FlashMLA nv_dev package.", +) + +_SUPPORTED_ATTENTION_HEAD_DIMS = (512, 576) +_VALUE_HEAD_DIM = 512 +_FLASH_MLA_TOPK_ALIGNMENT = 512 + + +def is_cudnn_sparse_attention_available() -> bool: + """Return whether the cuDNN backward and FlashMLA forward runtimes import.""" + return bool(_HAS_CUDNN_DSA and _HAS_FLASH_MLA) + + +def _require_available() -> None: + """Raise when either optional sparse-attention runtime is unavailable.""" + if not is_cudnn_sparse_attention_available(): + raise RuntimeError( + "cuDNN sparse attention requires both nvidia-cudnn-frontend[cutedsl] " + "and FlashMLA with flash_mla_sparse_fwd." + ) + + +def _require_cuda_tensors(operation: str, *tensors: torch.Tensor) -> tuple[int, int]: + """Validate that arbitrary-layout input tensors share one SM90+ CUDA device. + + Args: + operation: Name included in validation errors. + *tensors: Tensors with arbitrary shapes that must share one CUDA device. + + Returns: + CUDA compute capability as ``(major, minor)``. + """ + if not tensors or any(not tensor.is_cuda for tensor in tensors): + raise RuntimeError(f"{operation} requires CUDA tensors.") + device = tensors[0].device + if any(tensor.device != device for tensor in tensors[1:]): + raise ValueError(f"{operation} requires every tensor on the same CUDA device.") + major, minor = torch.cuda.get_device_capability(device) + if major < 9: + raise RuntimeError(f"{operation} requires SM90 or later, got SM{major}{minor}.") + return major, minor + + +def _compact_and_sort_indices(indices: torch.Tensor, key_count: int) -> tuple[torch.Tensor, torch.Tensor]: + """Canonicalize sparse indices into an ascending valid prefix. + + Args: + indices: Integer tensor of shape ``[query_tokens, sparse_width]`` with + global K/V coordinates and negative invalid entries. + key_count: Number of rows in the flattened K/V tensor. + + Returns: + A contiguous int32 index tensor of shape ``[query_tokens, sparse_width]`` + with an ascending valid prefix and ``-1`` suffix, and a contiguous int32 + tensor of shape ``[query_tokens]`` containing valid-prefix lengths. + """ + valid = (indices >= 0) & (indices < key_count) + topk_length = valid.sum(dim=-1, dtype=torch.int32) + positions = torch.arange(indices.size(-1), device=indices.device).view(1, -1) + compact_order = torch.where(valid, positions, torch.full_like(positions, indices.size(-1))).argsort(dim=-1) + indices = torch.gather(indices, -1, compact_order) + compact_valid = torch.gather(valid, -1, compact_order) + indices = indices.masked_fill(~compact_valid, -1) + + prefix = positions < topk_length.unsqueeze(-1) + sort_key = torch.where(prefix, indices, torch.full_like(indices, key_count)) + sort_order = sort_key.argsort(dim=-1) + indices = torch.gather(indices, -1, sort_order) + sorted_valid = torch.gather(prefix.expand_as(indices), -1, sort_order) + return indices.masked_fill(~sorted_valid, -1).to(torch.int32).contiguous(), topk_length.contiguous() + + +def _padded_head_count(num_heads: int, major: int) -> int: + """Return the FlashMLA-supported query-head count for one SM generation.""" + if major >= 10: + for padded in (64, 128): + if num_heads == padded or (num_heads < padded and padded % num_heads == 0): + return padded + alignment = 128 + else: + alignment = 64 + if num_heads % alignment == 0: + return num_heads + if num_heads < alignment and alignment % num_heads == 0: + return alignment + raise ValueError(f"FlashMLA sparse prefill requires the query-head count to divide {alignment}, got H={num_heads}.") + + +def _pad_attention_heads( + q: torch.Tensor, attn_sink: torch.Tensor, padded_heads: int +) -> tuple[torch.Tensor, torch.Tensor]: + """Pad query and attention-sink head axes for FlashMLA. + + Args: + q: Query tensor of shape ``[query_tokens, heads, head_dim]``. + attn_sink: FP32 attention-sink tensor of shape ``[heads]``. + padded_heads: FlashMLA-compatible output head count. + + Returns: + Query tensor of shape ``[query_tokens, padded_heads, head_dim]`` and + attention-sink tensor of shape ``[padded_heads]``. Inputs are returned + unchanged when ``heads == padded_heads``. + """ + if q.shape[1] == padded_heads: + return q, attn_sink + q_padded = q.new_zeros((q.shape[0], padded_heads, q.shape[2])) + q_padded[:, : q.shape[1]] = q + sink_padded = attn_sink.new_full((padded_heads,), float("-inf")) + sink_padded[: q.shape[1]] = attn_sink + return q_padded, sink_padded + + +class _CudnnSparseAttention(torch.autograd.Function): + """Pair FlashMLA forward with cuDNN backward for latent THD attention.""" + + @staticmethod + def forward( + ctx: Any, + q: torch.Tensor, + kv_latent: torch.Tensor, + topk_indices: torch.Tensor, + softmax_scale: float, + padded_heads: int, + topk_length: torch.Tensor | None, + all_rows_nonempty: bool, + valid_row_indices: torch.Tensor | None, + ) -> torch.Tensor: + """Run FlashMLA forward and save tensors required by cuDNN backward. + + Args: + ctx: Autograd context used to save forward tensors and scalar metadata. + q: CUDA BF16 query tensor of shape ``[query_tokens, heads, head_dim]``. + kv_latent: CUDA BF16 latent K/V tensor of shape ``[key_tokens, 1, head_dim]``. + topk_indices: CUDA int32 tensor of shape ``[query_tokens, 1, sparse_width]`` + containing global K/V coordinates and invalid entries marked ``-1``. + softmax_scale: Scale applied to query-key scores. + padded_heads: FlashMLA-compatible padded head count. + topk_length: Optional int32 valid-prefix lengths of shape ``[query_tokens]``. + all_rows_nonempty: Whether every query has a positive valid prefix. + valid_row_indices: Optional int64 indices of nonempty queries with shape + ``[valid_query_tokens]``. + + Returns: + CUDA BF16 latent values of shape ``[query_tokens, heads, 512]``. + """ + kv = kv_latent.squeeze(1).contiguous() + if topk_length is None: + indices, topk_length = _compact_and_sort_indices(topk_indices.squeeze(1), kv.shape[0]) + else: + indices = topk_indices.squeeze(1).contiguous() + padded_topk = math.ceil(indices.shape[-1] / _FLASH_MLA_TOPK_ALIGNMENT) * _FLASH_MLA_TOPK_ALIGNMENT + if padded_topk != indices.shape[-1]: + indices = torch.nn.functional.pad(indices, (0, padded_topk - indices.shape[-1]), value=-1) + + attn_sink = torch.full((q.shape[1],), float("-inf"), dtype=torch.float32, device=q.device) + q_kernel, sink_kernel = _pad_attention_heads(q.contiguous(), attn_sink, padded_heads) + out_kernel, _max_logits, lse_kernel = _FLASH_MLA_SPARSE_FWD( + q_kernel, + kv.unsqueeze(1), + indices.unsqueeze(1), + softmax_scale, + d_v=_VALUE_HEAD_DIM, + attn_sink=sink_kernel, + topk_length=topk_length, + indexer_topk=0, + ) + out = out_kernel[:, : q.shape[1]].contiguous() + lse = lse_kernel[:, : q.shape[1]].contiguous() + if not all_rows_nonempty: + out.masked_fill_(topk_length.eq(0).view(-1, 1, 1), 0) + cached_valid_rows = ( + valid_row_indices if valid_row_indices is not None else torch.empty(0, dtype=torch.int64, device=q.device) + ) + ctx.save_for_backward(q, kv, out, lse, attn_sink, indices.clamp_min(0), topk_length, cached_valid_rows) + ctx.softmax_scale = softmax_scale + ctx.padded_heads = padded_heads + ctx.all_rows_nonempty = all_rows_nonempty + ctx.has_cached_valid_rows = valid_row_indices is not None + return out + + @staticmethod + def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor | None, ...]: + """Map output gradients to query and gathered latent-KV layouts. + + Args: + ctx: Autograd context populated by :meth:`forward`. + grad_output: CUDA BF16 output gradient of shape + ``[query_tokens, heads, 512]``. + + Returns: + Gradients for the eight forward inputs: query tensor of shape + ``[query_tokens, heads, head_dim]``, latent K/V tensor of shape + ``[key_tokens, 1, head_dim]``, then ``None`` for scalar and metadata + inputs. + """ + q, kv, out, lse, attn_sink, indices, topk_length, cached_valid_rows = ctx.saved_tensors + valid_row_indices = None + if not ctx.all_rows_nonempty: + valid_row_indices = ( + cached_valid_rows + if ctx.has_cached_valid_rows + else torch.nonzero(topk_length > 0, as_tuple=False).flatten() + ) + + q_input = q + out_input = out + grad_input = grad_output + lse_input = lse + indices_kernel = indices + topk_length_kernel = topk_length + used_dummy = False + if valid_row_indices is not None: + if valid_row_indices.numel() == 0: + used_dummy = True + q_input = torch.zeros_like(q[:1]) + out_input = torch.zeros_like(out[:1]) + grad_input = torch.zeros_like(grad_output[:1]) + lse_input = torch.zeros_like(lse[:1]) + indices_kernel = torch.zeros_like(indices[:1]) + topk_length_kernel = torch.ones_like(topk_length[:1]) + else: + q_input = q.index_select(0, valid_row_indices) + out_input = out.index_select(0, valid_row_indices) + grad_input = grad_output.index_select(0, valid_row_indices) + lse_input = lse.index_select(0, valid_row_indices) + indices_kernel = indices.index_select(0, valid_row_indices) + topk_length_kernel = topk_length.index_select(0, valid_row_indices) + + q_kernel, sink_kernel = _pad_attention_heads(q_input, attn_sink, ctx.padded_heads) + if ctx.padded_heads == q.shape[1]: + out_kernel = out_input + grad_kernel = grad_input.contiguous() + lse_kernel = lse_input + else: + out_kernel = out_input.new_zeros((out_input.shape[0], ctx.padded_heads, out_input.shape[2])) + out_kernel[:, : out_input.shape[1]] = out_input + grad_kernel = grad_input.new_zeros((grad_input.shape[0], ctx.padded_heads, grad_input.shape[2])) + grad_kernel[:, : grad_input.shape[1]] = grad_input + lse_kernel = lse_input.new_zeros((lse_input.shape[0], ctx.padded_heads)) + lse_kernel[:, : lse_input.shape[1]] = lse_input + + result = _CUDNN_DSA.sparse_attention_backward_wrapper( + q_kernel.contiguous(), + kv, + out_kernel.contiguous(), + grad_kernel.contiguous(), + lse_kernel.contiguous(), + sink_kernel, + indices_kernel, + softmax_scale=ctx.softmax_scale, + topk_length=topk_length_kernel, + ) + if valid_row_indices is None: + grad_q = result["dq"][:, : q.shape[1]].contiguous() + else: + grad_q_valid = result["dq"][:0, : q.shape[1]] if used_dummy else result["dq"][:, : q.shape[1]] + grad_q = torch.zeros_like(q) + grad_q.index_copy_(0, valid_row_indices, grad_q_valid) + grad_kv = result["dkv"].unsqueeze(1).contiguous() + return grad_q, grad_kv, None, None, None, None, None, None + + +def cudnn_sparse_attention( + q: torch.Tensor, + kv_latent: torch.Tensor, + topk_indices: torch.Tensor, + softmax_scale: float, + topk_length: torch.Tensor | None = None, + all_rows_nonempty: bool = False, + valid_row_indices: torch.Tensor | None = None, +) -> torch.Tensor: + """Run sparse latent attention with FlashMLA forward and cuDNN backward. + + Args: + q: Contiguous CUDA BF16 query tensor of shape + ``[query_tokens, heads, head_dim]``, where ``head_dim`` is 512 or 576. + kv_latent: Contiguous CUDA BF16 latent K/V tensor of shape + ``[key_tokens, 1, head_dim]``. + topk_indices: Contiguous CUDA int32 tensor of shape + ``[query_tokens, 1, sparse_width]`` with global K/V coordinates and + invalid entries marked ``-1``. + softmax_scale: Scale forwarded unchanged to FlashMLA and cuDNN backward. + topk_length: Optional contiguous CUDA int32 valid-prefix lengths of shape + ``[query_tokens]``. When supplied, ``topk_indices`` must already contain + a compact, ascending valid prefix. + all_rows_nonempty: Whether every query has a positive valid-prefix length. + valid_row_indices: Optional contiguous CUDA int64 indices of nonempty queries + with shape ``[valid_query_tokens]``. + + Returns: + Contiguous CUDA BF16 latent output tensor of shape + ``[query_tokens, heads, 512]``. + + Raises: + RuntimeError: If optional kernels, CUDA, or SM90+ are unavailable. + TypeError: If compute tensors are not BF16 or indices are not int32. + ValueError: If tensor layouts, dimensions, sparse width, or scale are invalid. + """ + _require_available() + major, _ = _require_cuda_tensors("cuDNN DSA sparse attention", q, kv_latent, topk_indices) + if q.dtype != torch.bfloat16 or kv_latent.dtype != torch.bfloat16: + raise TypeError(f"q and kv_latent must be bfloat16, got {q.dtype} and {kv_latent.dtype}.") + if topk_indices.dtype != torch.int32: + raise TypeError(f"topk_indices must be int32, got {topk_indices.dtype}.") + if q.ndim != 3 or q.shape[-1] not in _SUPPORTED_ATTENTION_HEAD_DIMS: + raise ValueError( + f"q must have shape [query_tokens, heads, head_dim] with head_dim in " + f"{_SUPPORTED_ATTENTION_HEAD_DIMS}, got {tuple(q.shape)}." + ) + head_dim = q.shape[-1] + if kv_latent.ndim != 3 or kv_latent.shape[1:] != (1, head_dim): + raise ValueError(f"kv_latent must have shape [key_tokens, 1, {head_dim}], got {tuple(kv_latent.shape)}.") + if topk_indices.ndim != 3 or topk_indices.shape[:2] != (q.shape[0], 1): + raise ValueError( + f"topk_indices must have shape [query_tokens, 1, sparse_width], got {tuple(topk_indices.shape)}." + ) + if topk_indices.shape[-1] <= 0: + raise ValueError(f"sparse_width must be positive, got {topk_indices.shape[-1]}.") + if kv_latent.shape[0] >= torch.iinfo(torch.int32).max: + raise ValueError("The flattened K/V token count must fit in an int32 global index.") + if not isinstance(softmax_scale, (float, int)) or not math.isfinite(float(softmax_scale)): + raise TypeError("softmax_scale must be a finite Python float.") + if float(softmax_scale) <= 0.0: + raise ValueError(f"softmax_scale must be positive, got {softmax_scale}.") + if not isinstance(all_rows_nonempty, bool): + raise TypeError(f"all_rows_nonempty must be a bool, got {type(all_rows_nonempty).__name__}.") + if all_rows_nonempty and valid_row_indices is not None: + raise ValueError("valid_row_indices must be None when all_rows_nonempty is true.") + if topk_length is not None: + if topk_length.shape != (q.shape[0],) or topk_length.dtype != torch.int32 or topk_length.device != q.device: + raise ValueError( + "topk_length must be an int32 tensor on the query device with shape " + f"{(q.shape[0],)}, got shape={tuple(topk_length.shape)}, " + f"dtype={topk_length.dtype}, device={topk_length.device}." + ) + if not topk_length.is_contiguous(): + raise ValueError("topk_length must be contiguous.") + if valid_row_indices is not None: + if topk_length is None: + raise ValueError("valid_row_indices requires precomputed topk_length metadata.") + if ( + valid_row_indices.ndim != 1 + or valid_row_indices.dtype != torch.int64 + or valid_row_indices.device != q.device + or not valid_row_indices.is_contiguous() + ): + raise ValueError("valid_row_indices must be a contiguous int64 tensor on the query device.") + if valid_row_indices.numel() > q.shape[0]: + raise ValueError("valid_row_indices cannot contain more entries than query rows.") + + padded_heads = _padded_head_count(q.shape[1], major) + return _CudnnSparseAttention.apply( + q.contiguous(), + kv_latent.contiguous(), + topk_indices.contiguous(), + float(softmax_scale), + padded_heads, + topk_length, + all_rows_nonempty, + valid_row_indices, + ) + + +__all__ = ["cudnn_sparse_attention", "is_cudnn_sparse_attention_available"] diff --git a/nemo_automodel/components/models/glm5_next/__init__.py b/nemo_automodel/components/models/glm5_next/__init__.py new file mode 100644 index 0000000000..3af1aec9bf --- /dev/null +++ b/nemo_automodel/components/models/glm5_next/__init__.py @@ -0,0 +1,20 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native AutoModel support for ``zai-org/GLM-5.3-Flash``.""" + +from nemo_automodel.components.models.glm5_next.config import ( + Glm5NextConfig, + Glm5NextTextConfig, + Glm5NextVisionConfig, +) +from nemo_automodel.components.models.glm5_next.model import Glm5NextForConditionalGeneration + +ModelClass = Glm5NextForConditionalGeneration + +__all__ = [ + "Glm5NextConfig", + "Glm5NextForConditionalGeneration", + "Glm5NextTextConfig", + "Glm5NextVisionConfig", +] diff --git a/nemo_automodel/components/models/glm5_next/config.py b/nemo_automodel/components/models/glm5_next/config.py new file mode 100644 index 0000000000..660efee4f1 --- /dev/null +++ b/nemo_automodel/components/models/glm5_next/config.py @@ -0,0 +1,325 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""AutoModel-owned configuration for GLM-5.3-Flash. + +The released checkpoint requires Transformers 5.16, while AutoModel's current +runtime baseline predates the upstream ``glm5_next`` config. These classes keep +the checkpoint field protocol stable and allow ``AutoConfig`` to resolve the +model without remote code or a dependency bump. +""" + +from __future__ import annotations + +from typing import Any + +from transformers.configuration_utils import PretrainedConfig + + +def _json_safe_value(value: Any) -> Any: + """Return a JSON-serializable representation of a config value.""" + if value.__class__.__module__ == "torch" and value.__class__.__name__ == "dtype": + return str(value).removeprefix("torch.") + if isinstance(value, dict): + return {key: _json_safe_value(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_json_safe_value(item) for item in value] + return value + + +class Glm5NextTextConfig(PretrainedConfig): + """Configuration for the GLM-5.3 hybrid KDA/KPool-DSA text backbone.""" + + model_type = "glm5_next_text" + keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = {"num_local_experts": "n_routed_experts"} + + def __init__( + self, + vocab_size: int = 154880, + hidden_size: int = 4096, + intermediate_size: int = 12288, + moe_intermediate_size: int = 2048, + num_hidden_layers: int = 45, + num_attention_heads: int = 64, + num_key_value_heads: int = 64, + n_shared_experts: int = 1, + n_routed_experts: int = 288, + routed_scaling_factor: float = 2.5, + kv_lora_rank: int = 512, + q_lora_rank: int = 1536, + qk_rope_head_dim: int = 0, + qk_nope_head_dim: int = 256, + v_head_dim: int = 256, + n_group: int = 1, + topk_group: int = 1, + num_experts_per_tok: int = 8, + norm_topk_prob: bool = True, + mlp_layer_types: list[str] | None = None, + layer_types: list[str] | None = None, + indexer_types: list[str] | None = None, + index_topk_pattern: str | list[str] | None = None, + index_topk_freq: int = 1, + index_skip_topk_offset: int = 2, + index_topk: int = 2048, + index_head_dim: int = 128, + index_n_heads: int = 32, + index_kpool: int = 16, + index_kpool_always_select_tail: bool = True, + hidden_act: str = "silu", + swiglu_limit: float = 10.0, + linear_head_dim: int = 128, + linear_num_heads: int = 64, + linear_conv_kernel_dim: int = 4, + linear_lower_bound: float | None = -5.0, + linear_attn_config: dict[str, Any] | None = None, + hc_mult: int = 4, + hc_eps: float = 1e-6, + hc_sinkhorn_iters: int = 20, + max_position_embeddings: int = 1048576, + initializer_range: float = 0.02, + rms_norm_eps: float = 1e-5, + use_cache: bool = False, + attention_bias: bool = False, + attention_dropout: float = 0.0, + output_router_logits: bool = False, + router_aux_loss_coef: float = 0.001, + num_nextn_predict_layers: int = 1, + pad_token_id: int | None = 154820, + bos_token_id: int | None = None, + eos_token_id: int | list[int] | None = None, + tie_word_embeddings: bool = False, + **kwargs: Any, + ) -> None: + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.n_shared_experts = n_shared_experts + self.n_routed_experts = n_routed_experts + self.routed_scaling_factor = routed_scaling_factor + self.kv_lora_rank = kv_lora_rank + self.q_lora_rank = q_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + # Transformers aliases ``head_dim`` to the RoPE-only width for DSA. + self.head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.n_group = n_group + self.topk_group = topk_group + self.num_experts_per_tok = num_experts_per_tok + self.norm_topk_prob = norm_topk_prob + + if mlp_layer_types is None: + dense = min(3, num_hidden_layers) + mlp_layer_types = ["dense"] * dense + ["sparse"] * (num_hidden_layers - dense) + if len(mlp_layer_types) != num_hidden_layers: + raise ValueError("mlp_layer_types must have one entry per decoder layer") + self.mlp_layer_types = list(mlp_layer_types) + + if layer_types is None: + layer_types = [ + "deepseek_sparse_attention" if layer_idx % 4 == 3 else "linear_attention" + for layer_idx in range(num_hidden_layers) + ] + layer_types = [ + "deepseek_sparse_attention" if layer_type == "full_attention" else layer_type for layer_type in layer_types + ] + if len(layer_types) != num_hidden_layers: + raise ValueError("layer_types must have one entry per decoder layer") + unknown_layer_types = set(layer_types) - {"linear_attention", "deepseek_sparse_attention"} + if unknown_layer_types: + raise ValueError(f"Unsupported GLM-5.3 attention layer types: {sorted(unknown_layer_types)}") + self.layer_types = list(layer_types) + + if indexer_types is None: + if index_topk_pattern is not None: + indexer_types = ( + [{"F": "full", "S": "shared"}[char] for char in index_topk_pattern] + if isinstance(index_topk_pattern, str) + else list(index_topk_pattern) + ) + else: + freq = max(int(index_topk_freq), 1) + indexer_types = [ + "full" if max(layer_idx - index_skip_topk_offset + 1, 0) % freq == 0 else "shared" + for layer_idx in range(num_hidden_layers) + ] + if len(indexer_types) != num_hidden_layers: + raise ValueError("indexer_types must have one entry per decoder layer") + self.indexer_types = list(indexer_types) + self.index_topk_pattern = index_topk_pattern + self.index_topk_freq = index_topk_freq + self.index_skip_topk_offset = index_skip_topk_offset + self.index_topk = index_topk + self.index_head_dim = index_head_dim + self.index_n_heads = index_n_heads + self.index_kpool = index_kpool + self.index_kpool_always_select_tail = index_kpool_always_select_tail + + if linear_attn_config is not None: + linear_head_dim = linear_attn_config.get("head_dim", linear_head_dim) + linear_num_heads = linear_attn_config.get("num_heads", linear_num_heads) + linear_conv_kernel_dim = linear_attn_config.get("short_conv_kernel_size", linear_conv_kernel_dim) + linear_lower_bound = linear_attn_config.get("gate_lower_bound", linear_lower_bound) + if linear_attn_config.get("safe_gate", True) and linear_lower_bound is None: + linear_lower_bound = -5.0 + self.linear_head_dim = linear_head_dim + self.linear_num_heads = linear_num_heads + self.linear_conv_kernel_dim = linear_conv_kernel_dim + self.linear_lower_bound = linear_lower_bound + self.linear_attn_config = { + "head_dim": linear_head_dim, + "num_heads": linear_num_heads, + "short_conv_kernel_size": linear_conv_kernel_dim, + "gate_lower_bound": linear_lower_bound, + "safe_gate": linear_lower_bound is not None, + } + + self.hc_mult = hc_mult + self.hc_eps = hc_eps + self.hc_sinkhorn_iters = hc_sinkhorn_iters + self.hidden_act = hidden_act + self.swiglu_limit = swiglu_limit + self.max_position_embeddings = max_position_embeddings + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.output_router_logits = output_router_logits + self.router_aux_loss_coef = router_aux_loss_coef + self.num_nextn_predict_layers = num_nextn_predict_layers + + if num_attention_heads != num_key_value_heads: + raise ValueError("GLM-5.3 DSA requires num_attention_heads == num_key_value_heads") + if q_lora_rank is None: + raise ValueError("GLM-5.3 DSA requires q_lora_rank") + if qk_rope_head_dim != 0: + raise ValueError("GLM-5.3 DSA is NoPE and requires qk_rope_head_dim=0") + if index_kpool < 1 or index_topk % index_kpool: + raise ValueError("index_kpool must be positive and divide index_topk") + + # ``head_dim`` is independently emitted by the checkpoint and must not + # overwrite the NoPE width above through a future attribute alias. + kwargs.pop("head_dim", None) + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + @property + def num_local_experts(self) -> int: + """Alias used by Hugging Face expert implementations.""" + return self.n_routed_experts + + def to_dict(self) -> dict[str, Any]: + return _json_safe_value(super().to_dict()) + + +class Glm5NextVisionConfig(PretrainedConfig): + """Configuration for the GLM-5.3 image/video encoder.""" + + model_type = "glm5_next_vision" + + def __init__( + self, + depth: int = 24, + hidden_size: int = 1024, + hidden_act: str = "silu", + attention_bias: bool = True, + attention_dropout: float = 0.0, + num_heads: int = 16, + in_channels: int = 3, + image_size: int = 448, + patch_size: int = 14, + rms_norm_eps: float = 1e-5, + spatial_merge_size: int = 2, + temporal_patch_size: int = 2, + out_hidden_size: int = 4096, + intermediate_size: int = 4096, + projection_intermediate_size: int = 10240, + initializer_range: float = 0.02, + swiglu_limit: float = 10.0, + **kwargs: Any, + ) -> None: + self.depth = depth + self.num_hidden_layers = depth + self.hidden_size = hidden_size + self.hidden_act = hidden_act + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.num_heads = num_heads + self.num_attention_heads = num_heads + self.in_channels = in_channels + self.image_size = image_size + self.patch_size = patch_size + self.rms_norm_eps = rms_norm_eps + self.spatial_merge_size = spatial_merge_size + self.temporal_patch_size = temporal_patch_size + self.out_hidden_size = out_hidden_size + self.intermediate_size = intermediate_size + self.projection_intermediate_size = projection_intermediate_size + self.initializer_range = initializer_range + self.swiglu_limit = swiglu_limit + super().__init__(**kwargs) + + def to_dict(self) -> dict[str, Any]: + return _json_safe_value(super().to_dict()) + + +class Glm5NextConfig(PretrainedConfig): + """Top-level GLM-5.3 vision-language configuration.""" + + model_type = "glm5_next" + sub_configs = {"text_config": Glm5NextTextConfig, "vision_config": Glm5NextVisionConfig} + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + text_config: dict[str, Any] | Glm5NextTextConfig | None = None, + vision_config: dict[str, Any] | Glm5NextVisionConfig | None = None, + image_token_id: int = 154854, + video_token_id: int = 154855, + image_start_token_id: int = 154830, + image_end_token_id: int = 154831, + video_start_token_id: int = 154832, + video_end_token_id: int = 154833, + tie_word_embeddings: bool = False, + **kwargs: Any, + ) -> None: + if text_config is None: + text_config = Glm5NextTextConfig() + elif isinstance(text_config, dict): + text_config = Glm5NextTextConfig(**text_config) + if vision_config is None: + vision_config = Glm5NextVisionConfig() + elif isinstance(vision_config, dict): + vision_config = Glm5NextVisionConfig(**vision_config) + self.text_config = text_config + self.vision_config = vision_config + self.image_token_id = image_token_id + self.video_token_id = video_token_id + self.image_start_token_id = image_start_token_id + self.image_end_token_id = image_end_token_id + self.video_start_token_id = video_start_token_id + self.video_end_token_id = video_end_token_id + self.hidden_size = text_config.hidden_size + self.vocab_size = text_config.vocab_size + self.max_position_embeddings = text_config.max_position_embeddings + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) + + def get_text_config(self, decoder: bool = False) -> Glm5NextTextConfig: + """Return the decoder config using the Transformers multimodal protocol.""" + del decoder + return self.text_config + + def to_dict(self) -> dict[str, Any]: + return _json_safe_value(super().to_dict()) diff --git a/nemo_automodel/components/models/glm5_next/cp.py b/nemo_automodel/components/models/glm5_next/cp.py new file mode 100644 index 0000000000..99299ffe90 --- /dev/null +++ b/nemo_automodel/components/models/glm5_next/cp.py @@ -0,0 +1,341 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Contiguous packed context parallelism for GLM-5.3-Flash. + +Kimi Delta Attention carries recurrent and short-convolution state from left to +right, so the generic load-balanced CP permutation is not valid. This module +keeps one contiguous token interval per CP rank and one global document-id map +that both KDA and KPool-DSA consume. +""" + +from __future__ import annotations + +import contextlib +from dataclasses import dataclass, field +from typing import Any + +import torch +import torch.distributed as dist + +from nemo_automodel.components.distributed.context_parallel.sharder import ShardLayout +from nemo_automodel.shared.import_utils import safe_import_from + +_PAD_DOC_ID = 0 +_FLA_CP_AVAILABLE, _build_cp_context = safe_import_from( + "fla.ops.cp", + "build_cp_context", + msg="GLM-5.3 context parallelism requires the `fla` optional dependency.", +) + + +@dataclass +class Glm5NextPackedContext: + """Global packed-document layout for one model step. + + Attributes: + doc_ids: Integer document ids ``[batch, global_sequence]``. Zero marks + padding and positive values identify independent packed documents. + seq_start: Global offset of this rank's contiguous local token interval. + cp_size: Number of context-parallel ranks. + original_seq_len: Sequence length before CP divisibility padding. + """ + + doc_ids: torch.Tensor + seq_start: int = 0 + cp_size: int = 1 + original_seq_len: int | None = None + # FSDP mixed-precision hooks recursively rebuild dataclass kwargs with + # ``dataclasses.replace``. Keep the cache constructor-visible so that + # transform remains valid after packed metadata enters an FSDP root. + _cu_seqlens: dict[int, tuple[torch.Tensor, torch.Tensor]] = field( + default_factory=dict, + repr=False, + compare=False, + ) + + @property + def cp_enabled(self) -> bool: + """Return whether the sequence is split across more than one rank.""" + return self.cp_size > 1 + + @property + def local_seq_len(self) -> int: + """Return the padded local sequence length.""" + return self.doc_ids.shape[1] // self.cp_size + + @property + def local_doc_ids(self) -> torch.Tensor: + """Return document ids ``[batch, local_sequence]`` for this rank.""" + return self.doc_ids[:, self.seq_start : self.seq_start + self.local_seq_len] + + def row_cu_seqlens(self, row: int) -> tuple[torch.Tensor, torch.Tensor]: + """Return device/CPU segment boundaries for one packed batch row.""" + cached = self._cu_seqlens.get(row) + if cached is None: + device_boundaries = segment_cu_seqlens(self.doc_ids[row]).to(torch.long) + cached = (device_boundaries, device_boundaries.cpu()) + self._cu_seqlens[row] = cached + return cached + + +def doc_ids_from_seq_lens(seq_lens: torch.Tensor, seq_len: int, *, padding_value: int = -1000) -> torch.Tensor: + """Convert per-document lengths ``[batch, documents]`` to ids ``[batch, sequence]``.""" + if seq_lens.ndim == 1: + seq_lens = seq_lens.unsqueeze(0) + doc_ids = torch.zeros((seq_lens.shape[0], seq_len), dtype=torch.int32, device=seq_lens.device) + for row in range(seq_lens.shape[0]): + offset = 0 + lengths = seq_lens[row] + lengths = lengths[lengths != padding_value] + for doc_index, length in enumerate(lengths.tolist()): + if length <= 0 or offset >= seq_len: + continue + end = min(offset + int(length), seq_len) + doc_ids[row, offset:end] = doc_index + 1 + offset = end + return doc_ids + + +def doc_ids_from_cu_seqlens(cu_seqlens: torch.Tensor, seq_len: int) -> torch.Tensor: + """Convert cumulative document boundaries to ids ``[1, sequence]``.""" + boundaries = cu_seqlens.flatten().to(torch.long) + boundaries = boundaries[boundaries >= 0] + if boundaries.numel() < 2: + return torch.ones((1, seq_len), dtype=torch.int32, device=cu_seqlens.device) + positions = torch.arange(seq_len, device=cu_seqlens.device) + doc_ids = torch.bucketize(positions, boundaries[1:], right=True) + 1 + doc_ids = torch.where(positions < boundaries[-1], doc_ids, torch.zeros_like(doc_ids)) + return doc_ids.to(torch.int32).unsqueeze(0) + + +def segment_cu_seqlens(doc_ids_row: torch.Tensor) -> torch.Tensor: + """Return boundaries for consecutive document-id runs covering the full row.""" + if doc_ids_row.numel() == 0: + return torch.zeros(1, dtype=torch.int32, device=doc_ids_row.device) + starts = torch.nonzero(doc_ids_row[1:] != doc_ids_row[:-1], as_tuple=False).flatten() + 1 + return torch.cat( + ( + torch.zeros(1, dtype=starts.dtype, device=starts.device), + starts, + torch.full((1,), doc_ids_row.numel(), dtype=starts.dtype, device=starts.device), + ) + ).to(torch.int32) + + +def build_fla_cp_context( + packed_context: Glm5NextPackedContext, + row: int, + cp_group: Any, + conv_kernel_size: int, +): + """Build FLA's KDA context for one batch row. + + Args: + packed_context: Global document layout. + row: Batch row being executed. + cp_group: Context-parallel process group. + conv_kernel_size: Short-convolution width used for the left halo. + + Returns: + FLA ``FLACPContext`` carrying segment and process-group metadata. + """ + if not _FLA_CP_AVAILABLE: + raise RuntimeError("GLM-5.3 context parallelism requires the `fla` optional dependency") + + cu_seqlens, cu_seqlens_cpu = packed_context.row_cu_seqlens(row) + return _build_cp_context( + cu_seqlens=cu_seqlens, + group=cp_group, + conv1d_kernel_size=conv_kernel_size, + cu_seqlens_cpu=cu_seqlens_cpu, + ) + + +class _AllGatherSequence(torch.autograd.Function): + """Autograd-aware all-gather of equal contiguous sequence shards.""" + + @staticmethod + def forward(ctx, local_tensor: torch.Tensor, group: Any, dim: int) -> torch.Tensor: + dim = dim if dim >= 0 else local_tensor.ndim + dim + gathered = [torch.empty_like(local_tensor) for _ in range(dist.get_world_size(group))] + dist.all_gather(gathered, local_tensor.contiguous(), group=group) + ctx.group = group + ctx.dim = dim + ctx.rank = dist.get_rank(group) + ctx.local_size = local_tensor.shape[dim] + return torch.cat(gathered, dim=dim) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + grad_output = grad_output.contiguous() + dist.all_reduce(grad_output, op=dist.ReduceOp.SUM, group=ctx.group) + start = ctx.rank * ctx.local_size + return grad_output.narrow(ctx.dim, start, ctx.local_size).contiguous(), None, None + + +class _AllGatherBackwardAnchor(torch.autograd.Function): + """Return zero while retaining a backward edge to a gathered tensor.""" + + @staticmethod + def forward(ctx, gathered: torch.Tensor) -> torch.Tensor: + ctx.input_shape = gathered.shape + return gathered.new_zeros((gathered.shape[0], 1, 1)) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor: + return grad_output.new_zeros(ctx.input_shape) + + +def all_gather_sequence(tensor: torch.Tensor, cp_group: Any, *, dim: int = 1) -> torch.Tensor: + """Gather a sequence-sharded tensor while preserving K/V gradient flow.""" + return _AllGatherSequence.apply(tensor, cp_group, dim) + + +def all_gather_backward_anchor(gathered: torch.Tensor) -> torch.Tensor: + """Create a zero-valued dependency that keeps gather backward collective-safe.""" + return _AllGatherBackwardAnchor.apply(gathered) + + +def _pad_sequence_dim(tensor: torch.Tensor, pad_len: int, value: float | int | bool) -> torch.Tensor: + if pad_len <= 0: + return tensor + pad = torch.full( + (tensor.shape[0], pad_len, *tensor.shape[2:]), + value, + dtype=tensor.dtype, + device=tensor.device, + ) + return torch.cat((tensor, pad), dim=1) + + +def _normalize_batch_axis(batch: dict[str, Any]) -> None: + """Restore the placeholder batch axis used by THD VLM collaters.""" + input_ids = batch["input_ids"] + if input_ids.ndim != 1: + return + sequence_length = input_ids.shape[0] + # Media tensors are flattened over vision patches, not text tokens. Restrict + # this repair to fields whose leading dimension is contractually the text + # sequence so an image with the same number of patches is never reshaped. + for key in ("input_ids", "labels", "position_ids", "attention_mask", "padding_mask", "_packed_seq_ids"): + value = batch.get(key) + if isinstance(value, torch.Tensor) and value.ndim >= 1 and value.shape[0] == sequence_length: + batch[key] = value.unsqueeze(0) + + +def _global_doc_ids_from_batch(batch: dict[str, Any], seq_len: int) -> torch.Tensor: + """Resolve the global document map before removing packing metadata.""" + packed_ids = batch.get("_packed_seq_ids") + if isinstance(packed_ids, torch.Tensor): + return packed_ids.to(torch.int32) if packed_ids.ndim == 2 else packed_ids.unsqueeze(0).to(torch.int32) + + attention_mask = batch.get("attention_mask") + if isinstance(attention_mask, torch.Tensor) and attention_mask.ndim == 2: + # Binary masks naturally describe one document; indexed masks already + # carry one-based packed document ids. + return attention_mask.to(torch.int32) + + seq_lens = batch.get("seq_lens_padded", batch.get("seq_lens")) + if isinstance(seq_lens, torch.Tensor): + return doc_ids_from_seq_lens(seq_lens, seq_len) + + cu_seqlens = batch.get("cu_seqlens") + if isinstance(cu_seqlens, torch.Tensor): + return doc_ids_from_cu_seqlens(cu_seqlens, seq_len) + + doc_ids = torch.ones( + (batch["input_ids"].shape[0], seq_len), + dtype=torch.int32, + device=batch["input_ids"].device, + ) + padding_mask = batch.get("padding_mask") + if isinstance(padding_mask, torch.Tensor): + doc_ids.masked_fill_(padding_mask.bool(), _PAD_DOC_ID) + return doc_ids + + +def shard_batch_for_glm5_next_cp( + cp_mesh, + tp_mesh, + batch: dict[str, Any], + *, + loss_mask=None, + padding_token_id: int = 0, + shard_primary: bool = False, +): + """Contiguously shard packed GLM-5.3 token streams. + + The top-level VLM uses ``shard_primary=False``: image features must be + spliced into the full embedding sequence inside ``forward`` before that + differentiable primary stream is sliced. Labels and other no-grad token + streams are still sharded here. + + Args: + cp_mesh: One-dimensional CP mesh or ``None``. + tp_mesh: Unused tensor-parallel mesh, accepted by the sharder protocol. + batch: Batch containing token tensors with shape ``[batch, sequence]``. + loss_mask: Optional loss mask ``[batch, sequence]``. + padding_token_id: Fill value for padded token ids. + shard_primary: Whether to shard ``input_ids`` in this function. + + Returns: + Context factory, mutated local batch, and the global shard layout. + """ + del tp_mesh + _normalize_batch_axis(batch) + input_ids = batch["input_ids"] + seq_len = input_ids.shape[1] + cp_size = 1 if cp_mesh is None else cp_mesh.size() + doc_ids = _global_doc_ids_from_batch(batch, seq_len) + + for key in ( + "attention_mask", + "_packed_seq_ids", + "seq_lens", + "seq_lens_padded", + "cu_seqlens", + "cu_seqlens_padded", + "max_seqlen", + "qkv_format", + ): + batch.pop(key, None) + + if "position_ids" not in batch: + batch["position_ids"] = torch.arange(seq_len, device=input_ids.device).unsqueeze(0).expand_as(input_ids) + + pad_len = (-seq_len) % cp_size + padded_seq_len = seq_len + pad_len + if pad_len: + for key, pad_value in (("labels", -100), ("position_ids", 0), ("padding_mask", True)): + value = batch.get(key) + if isinstance(value, torch.Tensor) and value.ndim >= 2 and value.shape[1] == seq_len: + batch[key] = _pad_sequence_dim(value, pad_len, pad_value) + if shard_primary: + batch["input_ids"] = _pad_sequence_dim(batch["input_ids"], pad_len, padding_token_id) + doc_ids = _pad_sequence_dim(doc_ids, pad_len, _PAD_DOC_ID) + if isinstance(loss_mask, torch.Tensor): + loss_mask = _pad_sequence_dim(loss_mask, pad_len, 0) + + batch.setdefault("padding_mask", doc_ids <= _PAD_DOC_ID) + seq_start = 0 if cp_mesh is None else cp_mesh.get_local_rank() * (padded_seq_len // cp_size) + batch["glm5_next_packed_context"] = Glm5NextPackedContext( + doc_ids=doc_ids, + seq_start=seq_start, + cp_size=cp_size, + original_seq_len=seq_len, + ) + + local_seq_len = padded_seq_len // cp_size + seq_end = seq_start + local_seq_len + shard_keys = ["labels", "position_ids", "padding_mask"] + if shard_primary: + shard_keys.append("input_ids") + for key in shard_keys: + value = batch.get(key) + if isinstance(value, torch.Tensor) and value.ndim >= 2 and value.shape[1] == padded_seq_len: + batch[key] = value[:, seq_start:seq_end].contiguous() + if isinstance(loss_mask, torch.Tensor): + batch["loss_mask"] = loss_mask[:, seq_start:seq_end].contiguous() + + return contextlib.nullcontext, batch, ShardLayout(original_seq_len=seq_len, padded_seq_len=padded_seq_len) diff --git a/nemo_automodel/components/models/glm5_next/image_processing.py b/nemo_automodel/components/models/glm5_next/image_processing.py new file mode 100644 index 0000000000..ff375270f2 --- /dev/null +++ b/nemo_automodel/components/models/glm5_next/image_processing.py @@ -0,0 +1,234 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Image-only backport of the GLM-5.3 dynamic patch processor.""" + +from __future__ import annotations + +import math + +import torch +from transformers.image_processing_backends import TorchvisionBackend +from transformers.image_processing_utils import BatchFeature +from transformers.image_transforms import group_images_by_shape, reorder_images +from transformers.image_utils import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD, PILImageResampling, SizeDict +from transformers.processing_utils import ImagesKwargs + +from nemo_automodel.shared.import_utils import safe_import + +_, tvF = safe_import( + "torchvision.transforms.v2.functional", + msg="GLM-5.3 image preprocessing requires torchvision. Install the `vlm` extra.", +) + + +class Glm5NextImageProcessorKwargs(ImagesKwargs, total=False): + """Additional dynamic resize and patchification options.""" + + patch_size: int + temporal_patch_size: int + merge_size: int + patch_expand_factor: int + min_image_tokens: int + max_image_tokens: int + + +def smart_resize( + num_frames: int, + height: int, + width: int, + temporal_factor: int = 2, + factor: int = 28, + min_pixels: int = 16, + max_pixels: int = 8000, +) -> tuple[int, int]: + """Return an aligned H/W canvas inside the configured token budget.""" + pixels_per_token = temporal_factor * factor**2 + min_pixels *= pixels_per_token + max_pixels *= pixels_per_token + + def align(value: int) -> int: + return math.ceil(value / factor) * factor + + aligned_frames = max(temporal_factor, round(num_frames / temporal_factor) * temporal_factor) + target_height, target_width = align(height), align(width) + budget = aligned_frames * target_height * target_width + if budget < min_pixels: + scale = math.sqrt(min_pixels / (num_frames * height * width)) + target_height, target_width = align(max(1, math.ceil(height * scale))), align(max(1, math.ceil(width * scale))) + budget = aligned_frames * target_height * target_width + if budget > max_pixels: + if max_pixels < aligned_frames * factor**2: + raise ValueError(f"max pixel budget {max_pixels} cannot hold one aligned patch") + low, high = 1, height + target_height = target_width = factor + while low <= high: + content_height = (low + high) // 2 + content_width = max(1, math.floor(width * content_height / height)) + candidate_height, candidate_width = align(content_height), align(content_width) + if aligned_frames * candidate_height * candidate_width <= max_pixels: + target_height, target_width = candidate_height, candidate_width + low = content_height + 1 + else: + high = content_height - 1 + return target_height, target_width + + +class Glm5NextImageProcessor(TorchvisionBackend): + """Dynamically resize, normalize and flatten GLM-5.3 image patches.""" + + do_resize = True + resample = PILImageResampling.BICUBIC + size = {"longest_edge": 1} + default_to_square = False + do_rescale = True + rescale_factor = 1 / 255 + do_normalize = True + image_mean = OPENAI_CLIP_MEAN + image_std = OPENAI_CLIP_STD + do_convert_rgb = True + patch_size = 14 + temporal_patch_size = 2 + merge_size = 2 + patch_expand_factor = 1 + min_image_tokens = 16 + max_image_tokens = 8000 + valid_kwargs = Glm5NextImageProcessorKwargs + model_input_names = ["pixel_values", "image_grid_thw"] + + def resize( + self, + images: torch.Tensor, + resample, + factor: int, + temporal_factor: int, + min_image_tokens: int, + max_image_tokens: int, + **kwargs, + ) -> torch.Tensor: + """Aspect-preserving resize followed by right/bottom zero padding.""" + del kwargs + height, width = images.shape[-2:] + target_height, target_width = smart_resize( + temporal_factor, + height, + width, + temporal_factor, + factor, + min_image_tokens, + max_image_tokens, + ) + pixels_per_token = temporal_factor * factor**2 + scale = min(target_height / height, target_width / width) + if temporal_factor * height * width >= pixels_per_token * min_image_tokens: + scale = min(1.0, scale) + content_height = max(1, min(target_height, math.floor(height * scale))) + content_width = max(1, min(target_width, math.floor(width * scale))) + if (content_height, content_width) != (height, width): + images = super().resize(images, SizeDict(height=content_height, width=content_width), resample=resample) + return tvF.pad(images, [0, 0, target_width - content_width, target_height - content_height], fill=0) + + @staticmethod + def patchify( + images: torch.Tensor, + patch_size: int, + merge_size: int, + temporal_patch_size: int, + ) -> tuple[torch.Tensor, int, int]: + """Flatten block-major duplicated temporal patches.""" + batch, channels, height, width = images.shape + grid_h, grid_w = height // patch_size, width // patch_size + patches = images.reshape( + batch, + channels, + grid_h // merge_size, + merge_size, + patch_size, + grid_w // merge_size, + merge_size, + patch_size, + ).permute(0, 2, 5, 3, 6, 1, 4, 7) + patches = ( + patches.unsqueeze(6) + .expand(-1, -1, -1, -1, -1, -1, temporal_patch_size, -1, -1) + .reshape(batch, grid_h * grid_w, channels * temporal_patch_size * patch_size * patch_size) + ) + return patches, grid_h, grid_w + + def _preprocess( + self, + images: list[torch.Tensor], + do_resize: bool, + size: SizeDict, + resample, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean, + image_std, + patch_size: int, + temporal_patch_size: int, + merge_size: int, + patch_expand_factor: int, + min_image_tokens: int, + max_image_tokens: int, + disable_grouping: bool | None, + return_tensors, + **kwargs, + ) -> BatchFeature: + """Implement TorchvisionBackend's grouped preprocessing contract.""" + del size, kwargs + grouped, indices = group_images_by_shape(images, disable_grouping=disable_grouping) + resized = {} + for shape, stacked in grouped.items(): + if do_resize: + stacked = self.resize( + stacked, + resample, + patch_size * merge_size * patch_expand_factor, + temporal_patch_size, + min_image_tokens, + max_image_tokens, + ) + resized[shape] = stacked + images = reorder_images(resized, indices) + grouped, indices = group_images_by_shape(images, disable_grouping=disable_grouping) + processed, grids = {}, {} + for shape, stacked in grouped.items(): + stacked = self.rescale_and_normalize( + stacked, + do_rescale, + rescale_factor, + do_normalize, + image_mean, + image_std, + ) + patches, grid_h, grid_w = self.patchify(stacked, patch_size, merge_size, temporal_patch_size) + processed[shape] = patches + grids[shape] = [[1, grid_h, grid_w]] * len(stacked) + images = reorder_images(processed, indices) + image_grids = reorder_images(grids, indices) + pixel_values = images[0] if len(images) == 1 else torch.cat(images, dim=0) + return BatchFeature( + data={"pixel_values": pixel_values, "image_grid_thw": torch.tensor(image_grids)}, + tensor_type=return_tensors, + ) + + def get_number_of_image_patches(self, height: int, width: int, images_kwargs: dict | None = None) -> int: + """Return the number of unmerged vision patches for one source image.""" + values = images_kwargs or {} + patch_size = values.get("patch_size", self.patch_size) + merge_size = values.get("merge_size", self.merge_size) + target_h, target_w = smart_resize( + self.temporal_patch_size, + height, + width, + self.temporal_patch_size, + patch_size * merge_size, + values.get("min_image_tokens", self.min_image_tokens), + values.get("max_image_tokens", self.max_image_tokens), + ) + return (target_h // patch_size) * (target_w // patch_size) + + +__all__ = ["Glm5NextImageProcessor", "smart_resize"] diff --git a/nemo_automodel/components/models/glm5_next/layers.py b/nemo_automodel/components/models/glm5_next/layers.py new file mode 100644 index 0000000000..a3c1e79dee --- /dev/null +++ b/nemo_automodel/components/models/glm5_next/layers.py @@ -0,0 +1,894 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native GLM-5.3 text layers. + +The released model alternates Kimi Delta Attention (KDA) and GLM-style +KPool-compressed Dynamic Sparse Attention (DSA), with manifold-constrained +Hyper-Connections (mHC) around both sublayers. FLA owns the production KDA +kernel; small pure-Torch fallbacks keep CPU construction and unit tests useful. +""" + +from __future__ import annotations + +import math +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn + +from nemo_automodel.components.distributed.activation_checkpointing import unwrap_checkpoint_wrapper +from nemo_automodel.components.models.common import BackendConfig +from nemo_automodel.components.models.common.cudnn_sparse_attention import ( + cudnn_sparse_attention, + is_cudnn_sparse_attention_available, +) +from nemo_automodel.components.models.glm5_next.config import Glm5NextTextConfig +from nemo_automodel.components.models.glm5_next.cp import ( + Glm5NextPackedContext, + all_gather_backward_anchor, + all_gather_sequence, + build_fla_cp_context, +) +from nemo_automodel.components.moe.config import MoEConfig +from nemo_automodel.components.moe.layers import MLP, MoE +from nemo_automodel.shared.import_utils import safe_import_from +from nemo_automodel.shared.utils import dtype_from_str as get_dtype + +_FLA_MSG = "GLM-5.3 KDA requires the flash-linear-attention/fla extra for GPU training." +_SHORT_CONV_OK, _fla_causal_conv1d = safe_import_from("fla.modules.conv", "causal_conv1d", msg=_FLA_MSG) +_CHUNK_KDA_OK, _chunk_kda = safe_import_from("fla.ops.kda", "chunk_kda", msg=_FLA_MSG) +_RECURRENT_KDA_OK, _recurrent_kda = safe_import_from("fla.ops.kda", "fused_recurrent_kda", msg=_FLA_MSG) +_KDA_GATE_OK, _fused_kda_gate = safe_import_from("fla.ops.kda.gate", "fused_kda_gate", msg=_FLA_MSG) + + +class Glm5NextRMSNorm(nn.Module): + """RMSNorm with fp32 variance accumulation. + + Input and output have shape ``[batch, sequence, hidden]``; the leading axes + may be replaced by any token layout as long as hidden is last. + """ + + def __init__(self, hidden_size: int, eps: float, dtype: torch.dtype) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Normalize ``[..., hidden]`` and preserve the input dtype.""" + input_dtype = hidden_states.dtype + states = hidden_states.float() + states = states * torch.rsqrt(states.square().mean(-1, keepdim=True) + self.variance_epsilon) + return self.weight * states.to(input_dtype) + + def reset_parameters(self) -> None: + nn.init.ones_(self.weight) + + +class Glm5NextUnweightedRMSNorm(nn.Module): + """Parameter-free fp32 RMS normalization used by mHC.""" + + def __init__(self, eps: float) -> None: + super().__init__() + self.eps = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Normalize ``[..., hc_streams * hidden]`` without a learned weight.""" + return hidden_states * torch.rsqrt(hidden_states.float().square().mean(-1, keepdim=True) + self.eps).to( + hidden_states.dtype + ) + + +class Glm5NextHyperConnectionFp32Params(nn.Module): + """Own mHC parameters that must remain fp32 under FSDP mixed precision.""" + + def __init__(self, mix_size: int) -> None: + super().__init__() + self.base = nn.Parameter(torch.empty(mix_size, dtype=torch.float32)) + self.scale = nn.Parameter(torch.empty(3, dtype=torch.float32)) + + def forward( + self, + pre_w: torch.Tensor, + post_w: torch.Tensor, + comb_w: torch.Tensor, + hc: int, + eps: float, + sinkhorn_iters: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build FP32 mHC weights while this holder's FSDP unit is unsharded.""" + pre_b, post_b, comb_b = self.base.split([hc, hc, hc * hc]) + pre_scale, post_scale, comb_scale = self.scale.unbind(0) + pre = torch.sigmoid(pre_w * pre_scale + pre_b) + eps + post = 2 * torch.sigmoid(post_w * post_scale + post_b) + comb_logits = comb_w.view(*comb_w.shape[:-1], hc, hc) * comb_scale + comb_b.view(hc, hc) + comb = torch.softmax(comb_logits, dim=-1) + eps + comb = comb / (comb.sum(dim=-2, keepdim=True) + eps) + for _ in range(sinkhorn_iters - 1): + comb = comb / (comb.sum(dim=-1, keepdim=True) + eps) + comb = comb / (comb.sum(dim=-2, keepdim=True) + eps) + return pre, post, comb + + +class Glm5NextHyperConnection(nn.Module): + """Manifold-constrained mixer for ``hc_mult`` residual streams. + + ``hidden_streams`` is ``[batch, sequence, hc_mult, hidden]``. The returned + tensors are ``post [batch, sequence, hc_mult]``, ``comb [batch, sequence, + hc_mult, hc_mult]`` and ``collapsed [batch, sequence, hidden]``. + """ + + def __init__(self, config: Glm5NextTextConfig) -> None: + super().__init__() + self.hc_mult = config.hc_mult + self.hc_sinkhorn_iters = config.hc_sinkhorn_iters + self.hc_eps = config.hc_eps + self.input_norm = Glm5NextUnweightedRMSNorm(config.rms_norm_eps) + mix = (2 + self.hc_mult) * self.hc_mult + dtype = get_dtype(getattr(config, "torch_dtype", None), torch.bfloat16) + self.fn = nn.Parameter(torch.empty(mix, self.hc_mult * config.hidden_size, dtype=dtype)) + self._fp32_params = Glm5NextHyperConnectionFp32Params(mix) + + @property + def base(self) -> nn.Parameter: + """Expose the checkpoint's flat mHC base parameter.""" + return self._fp32_params.base + + @property + def scale(self) -> nn.Parameter: + """Expose the checkpoint's flat mHC scale parameter.""" + return self._fp32_params.scale + + def forward(self, hidden_streams: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build Sinkhorn mixing weights and collapse streams for one sublayer.""" + hc = self.hc_mult + flat = self.input_norm(hidden_streams.flatten(start_dim=2).float()) + pre_w, post_w, comb_w = F.linear(flat, self.fn.float()).split([hc, hc, hc * hc], dim=-1) + pre, post, comb = self._fp32_params( + pre_w, + post_w, + comb_w, + hc, + self.hc_eps, + self.hc_sinkhorn_iters, + ) + collapsed = (pre.unsqueeze(-1) * hidden_streams).sum(dim=2).to(hidden_streams.dtype) + return post, comb, collapsed + + @torch.no_grad() + def init_weights(self, buffer_device: torch.device, init_std: float) -> None: + """Initialize mHC parameters on ``buffer_device``.""" + with buffer_device: + nn.init.normal_(self.fn, mean=0.0, std=init_std) + self.base.zero_() + self.scale.fill_(1.0) + + +class _TorchShortConvolution(nn.Module): + """Depthwise causal Conv1d matching FLA ``ShortConvolution`` state keys.""" + + def __init__(self, hidden_size: int, kernel_size: int, dtype: torch.dtype) -> None: + super().__init__() + self.kernel_size = kernel_size + self.weight = nn.Parameter(torch.empty(hidden_size, 1, kernel_size, dtype=dtype)) + + def forward( + self, + x: torch.Tensor, + *, + cu_seqlens: torch.Tensor | None = None, + **kwargs: Any, + ) -> tuple[torch.Tensor, None]: + """Convolve ``[batch, sequence, channels]`` and reset at packed boundaries.""" + if _SHORT_CONV_OK and x.is_cuda: + return _fla_causal_conv1d( + x=x, + weight=self.weight.squeeze(1), + bias=None, + initial_state=kwargs.get("cache"), + output_final_state=kwargs.get("output_final_state", False), + activation="silu", + cu_seqlens=cu_seqlens, + cp_context=kwargs.get("cp_context"), + ) + if cu_seqlens is None: + y = F.conv1d(x.transpose(1, 2), self.weight, groups=x.shape[-1], padding=self.kernel_size - 1) + return F.silu(y[..., : x.shape[1]].transpose(1, 2)), None + output = torch.zeros_like(x) + boundaries = cu_seqlens.flatten().tolist() + for start, end in zip(boundaries[:-1], boundaries[1:]): + if end <= start: + continue + segment = x[:, start:end] + y = F.conv1d(segment.transpose(1, 2), self.weight, groups=x.shape[-1], padding=self.kernel_size - 1) + output[:, start:end] = F.silu(y[..., : end - start].transpose(1, 2)) + return output, None + + def reset_parameters(self) -> None: + nn.init.uniform_(self.weight, -0.01, 0.01) + + +class _TorchRMSNormGated(Glm5NextRMSNorm): + """CPU fallback for FLA's gated per-head RMSNorm.""" + + def forward(self, hidden_states: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + """Normalize ``[batch, sequence, heads, head_dim]`` then sigmoid-gate it.""" + return super().forward(hidden_states) * torch.sigmoid(gate.float()).to(hidden_states.dtype) + + +def _short_conv(hidden_size: int, kernel_size: int, dtype: torch.dtype) -> nn.Module: + return _TorchShortConvolution(hidden_size, kernel_size, dtype) + + +def _rms_norm_gated(hidden_size: int, eps: float, dtype: torch.dtype) -> nn.Module: + return _TorchRMSNormGated(hidden_size, eps, dtype) + + +class Glm5NextKDAFp32Params(nn.Module): + """Own recurrent-decay parameters that must remain fp32 under FSDP.""" + + def __init__(self, num_heads: int, projection_size: int) -> None: + super().__init__() + # Keep the native layout identical to the released checkpoint. Besides + # avoiding a state-dict reshape, this matters under FSDP: checkpoint + # planning sees a DTensor sharded along dimension zero and cannot remove + # that dimension before the parameter is materialized for forward. + self.A_log = nn.Parameter(torch.empty(num_heads, dtype=torch.float32)) + self.dt_bias = nn.Parameter(torch.empty(projection_size, dtype=torch.float32)) + + def forward(self, gate: torch.Tensor, head_dim: int, lower_bound: float | None) -> torch.Tensor: + """Return log-decay gates ``[batch, sequence, heads, head_dim]``.""" + gate = gate.reshape(*gate.shape[:-1], -1, head_dim) + if _KDA_GATE_OK and gate.is_cuda: + return _fused_kda_gate( + gate, + self.A_log.contiguous(), + dt_bias=self.dt_bias.contiguous(), + lower_bound=lower_bound, + ) + gate = gate.float() + self.dt_bias.view(1, 1, -1, head_dim) + decay = self.A_log.view(1, 1, -1, 1).exp() + return lower_bound * torch.sigmoid(decay * gate) if lower_bound is not None else -decay * F.softplus(gate) + + +def _torch_recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + cu_seqlens: torch.Tensor | None, +) -> torch.Tensor: + """Differentiable reference KDA for CPU/small tensors. + + All q/k/v/g tensors are ``[batch, sequence, heads, head_dim]`` and beta is + ``[batch, sequence, heads]``. Packed boundaries reset the recurrent state. + """ + batch, sequence, heads, head_dim = q.shape + output = torch.zeros_like(v) + boundaries = [0, sequence] if cu_seqlens is None else cu_seqlens.flatten().tolist() + for start, end in zip(boundaries[:-1], boundaries[1:]): + state = torch.zeros(batch, heads, head_dim, head_dim, dtype=torch.float32, device=q.device) + for token in range(start, end): + q_t = q[:, token].float() * (head_dim**-0.5) + k_t, v_t = k[:, token].float(), v[:, token].float() + state = state * g[:, token].exp().unsqueeze(-1) + prediction = torch.einsum("bhd,bhdv->bhv", k_t, state) + error = (v_t - prediction) * beta[:, token].float().unsqueeze(-1) + state = state + torch.einsum("bhd,bhv->bhdv", k_t, error) + output[:, token] = torch.einsum("bhd,bhdv->bhv", q_t, state).to(output.dtype) + return output + + +class Glm5NextLinearAttention(nn.Module): + """Kimi Delta Attention with released GLM-5.3 checkpoint parameter names.""" + + def __init__(self, config: Glm5NextTextConfig, layer_idx: int) -> None: + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + self.head_dim = config.linear_head_dim + self.num_heads = config.linear_num_heads + self.projection_size = self.head_dim * self.num_heads + self.conv_size = config.linear_conv_kernel_dim + dtype = get_dtype(getattr(config, "torch_dtype", None), torch.bfloat16) + self.q_proj = nn.Linear(config.hidden_size, self.projection_size, bias=False, dtype=dtype) + self.k_proj = nn.Linear(config.hidden_size, self.projection_size, bias=False, dtype=dtype) + self.v_proj = nn.Linear(config.hidden_size, self.projection_size, bias=False, dtype=dtype) + self.q_conv1d = _short_conv(self.projection_size, self.conv_size, dtype) + self.k_conv1d = _short_conv(self.projection_size, self.conv_size, dtype) + self.v_conv1d = _short_conv(self.projection_size, self.conv_size, dtype) + self._fp32_params = Glm5NextKDAFp32Params(self.num_heads, self.projection_size) + self.f_a_proj = nn.Linear(config.hidden_size, self.head_dim, bias=False, dtype=dtype) + self.f_b_proj = nn.Linear(self.head_dim, self.projection_size, bias=False, dtype=dtype) + self.b_proj = nn.Linear(config.hidden_size, self.num_heads, bias=False, dtype=dtype) + self.g_a_proj = nn.Linear(config.hidden_size, self.head_dim, bias=False, dtype=dtype) + self.g_b_proj = nn.Linear(self.head_dim, self.projection_size, bias=False, dtype=dtype) + self.o_norm = _rms_norm_gated(self.head_dim, config.rms_norm_eps, dtype) + self.o_proj = nn.Linear(self.projection_size, config.hidden_size, bias=False, dtype=dtype) + self._cp_mesh = None + + @property + def A_log(self) -> nn.Parameter: + """Expose the checkpoint's flat ``A_log`` parameter name.""" + return self._fp32_params.A_log + + @property + def dt_bias(self) -> nn.Parameter: + """Expose the checkpoint's flat ``dt_bias`` parameter name.""" + return self._fp32_params.dt_bias + + def setup_cp_attention(self, cp_mesh) -> None: + """Attach the one-dimensional contiguous CP mesh.""" + self._cp_mesh = cp_mesh + + def forward( + self, + hidden_states: torch.Tensor, + *, + packed_context: Glm5NextPackedContext | None = None, + padding_mask: torch.Tensor | None = None, + **_: Any, + ) -> torch.Tensor: + """Run KDA over ``[batch, local_sequence, hidden]`` without crossing documents.""" + if packed_context is not None and packed_context.cp_enabled: + if self._cp_mesh is None: + raise RuntimeError("GLM-5.3 KDA received a CP batch before apply_cp attached its mesh") + group = self._cp_mesh.get_group() + outputs = [ + self._core( + hidden_states[row : row + 1], + cp_context=build_fla_cp_context(packed_context, row, group, self.conv_size), + ) + for row in range(hidden_states.shape[0]) + ] + output = torch.cat(outputs, dim=0) + elif packed_context is not None: + outputs = [] + for row in range(hidden_states.shape[0]): + cu_seqlens, _ = packed_context.row_cu_seqlens(row) + outputs.append(self._core(hidden_states[row : row + 1], cu_seqlens=cu_seqlens)) + output = torch.cat(outputs, dim=0) + else: + output = self._core(hidden_states) + if padding_mask is not None: + output = output.masked_fill(padding_mask.unsqueeze(-1), 0) + return output + + def _core( + self, + hidden_states: torch.Tensor, + *, + cu_seqlens: torch.Tensor | None = None, + cp_context: Any = None, + ) -> torch.Tensor: + """Project and execute KDA for one packed row or a regular batch.""" + kernel_kwargs: dict[str, Any] = {} if cp_context is None else {"cp_context": cp_context} + conv_kwargs = dict(cache=None, output_final_state=False, cu_seqlens=cu_seqlens, **kernel_kwargs) + q, _ = self.q_conv1d(x=self.q_proj(hidden_states), **conv_kwargs) + k, _ = self.k_conv1d(x=self.k_proj(hidden_states), **conv_kwargs) + v, _ = self.v_conv1d(x=self.v_proj(hidden_states), **conv_kwargs) + shape = (*hidden_states.shape[:-1], self.num_heads, self.head_dim) + q, k, v = q.view(shape).contiguous(), k.view(shape).contiguous(), v.view(shape).contiguous() + gate = self.f_b_proj(self.f_a_proj(hidden_states)).contiguous() + gate = self._fp32_params(gate, self.head_dim, self.config.linear_lower_bound).contiguous() + beta = self.b_proj(hidden_states).float().sigmoid().contiguous() + if _CHUNK_KDA_OK and hidden_states.is_cuda: + kernel = _chunk_kda if cp_context is not None or hidden_states.shape[1] > 64 else _recurrent_kda + kernel_options: dict[str, Any] = { + "use_qk_l2norm_in_kernel": True, + "transpose_state_layout": True, + } + if kernel is _chunk_kda: + kernel_options["safe_gate"] = self.config.linear_lower_bound is not None + output, _ = kernel( + q=q, + k=k, + v=v, + g=gate, + beta=beta, + initial_state=None, + output_final_state=cp_context is None, + cu_seqlens=cu_seqlens, + **kernel_options, + **kernel_kwargs, + ) + else: + q = (q.float() / torch.sqrt(q.float().square().sum(-1, keepdim=True) + 1e-6)).to(q.dtype) + k = (k.float() / torch.sqrt(k.float().square().sum(-1, keepdim=True) + 1e-6)).to(k.dtype) + output = _torch_recurrent_kda(q, k, v, gate, beta, cu_seqlens) + final_gate = self.g_b_proj(self.g_a_proj(hidden_states)).view(shape) + output = self.o_norm(output, final_gate).reshape(*hidden_states.shape[:-1], -1).contiguous() + return self.o_proj(output) + + @torch.no_grad() + def init_weights(self, buffer_device: torch.device, init_std: float) -> None: + """Initialize KDA while preserving fp32 recurrent parameters.""" + with buffer_device: + if self.config.linear_lower_bound is not None: + self.A_log.zero_() + else: + self.A_log.uniform_(1, 16).log_() + self.dt_bias.uniform_(math.log(1e-3), math.log(1e-1)) + dt = self.dt_bias.exp().clamp_min(1e-4) + self.dt_bias.copy_(dt + torch.log(-torch.expm1(-dt))) + for module in ( + self.q_proj, + self.k_proj, + self.v_proj, + self.f_a_proj, + self.f_b_proj, + self.b_proj, + self.g_a_proj, + self.g_b_proj, + self.o_proj, + ): + nn.init.normal_(module.weight, mean=0.0, std=init_std) + for conv in (self.q_conv1d, self.k_conv1d, self.v_conv1d): + conv.reset_parameters() + if hasattr(self.o_norm, "reset_parameters"): + self.o_norm.reset_parameters() + + +class Glm5NextKPoolIndexer(nn.Module): + """KPool-compressed DSA indexer for training without a KV cache.""" + + def __init__(self, config: Glm5NextTextConfig, layer_idx: int, dtype: torch.dtype) -> None: + super().__init__() + self.layer_idx = layer_idx + self.n_heads = config.index_n_heads + self.head_dim = config.index_head_dim + self.index_topk = config.index_topk + self.index_kpool = config.index_kpool + self.always_select_tail = config.index_kpool_always_select_tail + self.wq_b = nn.Linear(config.q_lora_rank, self.n_heads * self.head_dim, bias=False, dtype=dtype) + self.wk = nn.Linear(config.hidden_size, self.head_dim, bias=False, dtype=dtype) + self.k_norm = nn.LayerNorm(self.head_dim, eps=1e-6, dtype=dtype) + self.weights_proj = nn.Linear(config.hidden_size, self.n_heads, bias=False, dtype=dtype) + self.index_kpool_compress_ape = nn.Parameter(torch.zeros(self.index_kpool, self.head_dim, dtype=dtype)) + self.index_kpool_compress_gate = nn.Parameter(torch.zeros(self.head_dim, config.hidden_size, dtype=dtype)) + self.softmax_scale = self.head_dim**-0.5 + + @torch.no_grad() + def prepare_pools(self, full_hidden: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Compress one document's keys into KPool candidates. + + Args: + full_hidden: Hidden states with shape ``[1, keys, hidden]`` for one + unpadded document. + + Returns: + Pool keys with shape ``[complete_pools, index_head_dim]`` and raw + token indices with shape ``[complete_pools, index_kpool]``. + """ + keys = self.k_norm(self.wk(full_hidden)).squeeze(0) + gates = F.linear(full_hidden.squeeze(0), self.index_kpool_compress_gate) + length = keys.shape[0] + complete_pools = length // self.index_kpool + if complete_pools: + width = complete_pools * self.index_kpool + grouped_keys = keys[:width].view(complete_pools, self.index_kpool, self.head_dim) + grouped_gates = gates[:width].view(complete_pools, self.index_kpool, self.head_dim) + logits = grouped_gates.float() + self.index_kpool_compress_ape.float().unsqueeze(0) + pool_keys = (logits.softmax(dim=1).to(keys.dtype) * grouped_keys).sum(dim=1) + pool_indices = torch.arange(width, device=keys.device).view(complete_pools, self.index_kpool) + else: + pool_keys = keys.new_empty((0, self.head_dim)) + pool_indices = torch.empty((0, self.index_kpool), dtype=torch.long, device=keys.device) + return pool_keys, pool_indices + + @torch.no_grad() + def select( + self, + query_hidden: torch.Tensor, + query_resid: torch.Tensor, + query_positions: torch.Tensor, + pool_keys: torch.Tensor, + pool_indices: torch.Tensor, + key_length: int, + ) -> torch.Tensor: + """Select raw key indices for one query chunk. + + Args: + query_hidden: Hidden states with shape ``[1, queries, hidden]``. + query_resid: Low-rank query states with shape + ``[1, queries, q_lora_rank]``. + query_positions: Document-local positions with shape ``[queries]``. + pool_keys: Prepared KPool keys with shape + ``[complete_pools, index_head_dim]``. + pool_indices: Prepared raw token indices with shape + ``[complete_pools, index_kpool]``. + key_length: Number of tokens in the unpadded document. + + Returns: + Int32 raw indices with shape + ``[1, queries, index_topk + index_kpool - 1]`` when tail selection + is enabled, otherwise ``[1, queries, index_topk]``. + """ + complete_pools = pool_keys.shape[0] + + queries = self.wq_b(query_resid).view(1, -1, self.n_heads, self.head_dim) + scores = torch.einsum("bqhd,pd->bqhp", queries.float(), pool_keys.float()) + scores = F.relu(scores * self.softmax_scale) + weights = self.weights_proj(query_hidden).float() * (self.n_heads**-0.5) + scores = torch.einsum("bqh,bqhp->bqp", weights, scores) + if complete_pools: + pool_end = pool_indices[:, -1] + visible = pool_end.view(1, 1, -1) <= query_positions.view(1, -1, 1) + scores = scores.masked_fill(~visible, torch.finfo(scores.dtype).min) + select_k = min(self.index_topk // self.index_kpool, complete_pools) + selected = scores.topk(select_k, dim=-1).indices + selected_valid = visible.expand_as(scores).gather(-1, selected) + raw = pool_indices[selected].flatten(-2) + raw = raw.masked_fill(~selected_valid.unsqueeze(-1).expand_as(pool_indices[selected]).flatten(-2), -1) + else: + raw = torch.empty((1, query_hidden.shape[1], 0), dtype=torch.long, device=query_hidden.device) + + output_width = self.index_topk + if self.always_select_tail and self.index_kpool > 1: + tail_count = (query_positions + 1).remainder(self.index_kpool) + tail_start = query_positions + 1 - tail_count + offsets = torch.arange(self.index_kpool - 1, device=query_hidden.device) + tail = tail_start[:, None] + offsets + tail = tail.masked_fill(offsets[None] >= tail_count[:, None], -1).unsqueeze(0) + raw = torch.cat((raw, tail), dim=-1) + output_width += self.index_kpool - 1 + return F.pad(raw, (0, max(output_width - raw.shape[-1], 0)), value=-1)[..., :output_width].to(torch.int32) + + @torch.no_grad() + def forward( + self, + full_hidden: torch.Tensor, + query_hidden: torch.Tensor, + query_resid: torch.Tensor, + query_positions: torch.Tensor, + ) -> torch.Tensor: + """Prepare one document and select indices for a query chunk. + + ``full_hidden`` is ``[1, keys, hidden]``; query tensors are ``[1, + queries, ...]`` and positions are document-local ``[queries]``. + Returns ``[1, queries, index_topk + kpool - 1]`` int32 indices. + """ + pool_keys, pool_indices = self.prepare_pools(full_hidden) + return self.select( + query_hidden, + query_resid, + query_positions, + pool_keys, + pool_indices, + full_hidden.shape[1], + ) + + @torch.no_grad() + def init_weights(self, buffer_device: torch.device, init_std: float) -> None: + """Initialize indexer projections and KPool parameters.""" + with buffer_device: + for module in (self.wq_b, self.wk, self.weights_proj): + nn.init.normal_(module.weight, mean=0.0, std=init_std) + self.k_norm.reset_parameters() + self.index_kpool_compress_ape.zero_() + self.index_kpool_compress_gate.fill_(1.0) + + +class Glm5NextSparseAttention(nn.Module): + """NoPE MLA whose visibility is selected by the GLM KPool indexer.""" + + query_chunk_size = 32 + + def __init__(self, config: Glm5NextTextConfig, layer_idx: int, backend: BackendConfig) -> None: + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.backend = backend + self.num_heads = config.num_attention_heads + self.q_lora_rank = config.q_lora_rank + self.qk_head_dim = config.qk_nope_head_dim + self.v_head_dim = config.v_head_dim + self.kv_lora_rank = config.kv_lora_rank + self.scaling = self.qk_head_dim**-0.5 + if backend.attn == "cudnn" and self.kv_lora_rank != 512: + raise ValueError(f"GLM-5.3 cuDNN sparse attention requires kv_lora_rank=512, got {self.kv_lora_rank}.") + if backend.attn == "cudnn" and config.attention_dropout != 0.0: + raise ValueError( + "GLM-5.3 cuDNN sparse attention does not support attention dropout; " + f"got attention_dropout={config.attention_dropout}." + ) + dtype = get_dtype(getattr(config, "torch_dtype", None), torch.bfloat16) + self.q_a_proj = nn.Linear(config.hidden_size, self.q_lora_rank, bias=config.attention_bias, dtype=dtype) + self.q_a_layernorm = Glm5NextRMSNorm(self.q_lora_rank, config.rms_norm_eps, dtype) + self.q_b_proj = nn.Linear(self.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False, dtype=dtype) + self.kv_a_proj_with_mqa = nn.Linear( + config.hidden_size, self.kv_lora_rank, bias=config.attention_bias, dtype=dtype + ) + self.kv_a_layernorm = Glm5NextRMSNorm(self.kv_lora_rank, config.rms_norm_eps, dtype) + self.kv_b_proj = nn.Linear( + self.kv_lora_rank, + self.num_heads * (self.qk_head_dim + self.v_head_dim), + bias=False, + dtype=dtype, + ) + self.o_proj = nn.Linear(self.num_heads * self.v_head_dim, config.hidden_size, bias=False, dtype=dtype) + self.indexer = Glm5NextKPoolIndexer(config, layer_idx, dtype) + self._cp_mesh = None + + def setup_cp_attention(self, cp_mesh) -> None: + """Attach the CP mesh used for differentiable full-sequence gathering.""" + self._cp_mesh = cp_mesh + + def forward( + self, + hidden_states: torch.Tensor, + *, + packed_context: Glm5NextPackedContext, + padding_mask: torch.Tensor | None = None, + **_: Any, + ) -> torch.Tensor: + """Run document-isolated sparse attention on ``[batch, local_sequence, hidden]``.""" + if packed_context is None: + raise ValueError("GLM-5.3 sparse attention requires a packed document context") + if packed_context.cp_enabled: + if self._cp_mesh is None: + raise RuntimeError("GLM-5.3 DSA received a CP batch before apply_cp attached its mesh") + full_hidden = all_gather_sequence(hidden_states, self._cp_mesh.get_group(), dim=1) + else: + full_hidden = hidden_states + output = torch.zeros_like(hidden_states) + local_start = packed_context.seq_start + local_end = local_start + hidden_states.shape[1] + for row in range(hidden_states.shape[0]): + doc_ids = packed_context.doc_ids[row] + starts = torch.nonzero(doc_ids[1:] != doc_ids[:-1], as_tuple=False).flatten().add(1).tolist() + boundaries = [0, *starts, doc_ids.numel()] + for doc_start, doc_end in zip(boundaries[:-1], boundaries[1:]): + if int(doc_ids[doc_start]) <= 0: + continue + query_start, query_end = max(doc_start, local_start), min(doc_end, local_end) + if query_end <= query_start: + continue + doc = full_hidden[row : row + 1, doc_start:doc_end] + local_query_start = query_start - doc_start + local_query_end = query_end - doc_start + doc_output = self._forward_document(doc, local_query_start, local_query_end) + out_start = query_start - local_start + output[row : row + 1, out_start : out_start + doc_output.shape[1]] = doc_output + if packed_context.cp_enabled: + # A short packed sample can leave this contiguous CP interval with + # no valid queries. Keep the differentiable all-gather connected to + # the local output so every CP rank launches its backward AllReduce. + output = output + all_gather_backward_anchor(full_hidden) + if padding_mask is not None: + output = output.masked_fill(padding_mask.unsqueeze(-1), 0) + return output + + def _forward_document(self, full_hidden: torch.Tensor, query_start: int, query_end: int) -> torch.Tensor: + """Execute sparse attention for a local query interval of one full document. + + Args: + full_hidden: Unpadded document states with shape ``[1, key_tokens, hidden]``. + query_start: Inclusive document-local index of the first local query. + query_end: Exclusive document-local index of the final local query. + + Returns: + Projected attention output with shape + ``[1, query_end - query_start, hidden]``. + """ + length = full_hidden.shape[1] + latent = self.kv_a_layernorm(self.kv_a_proj_with_mqa(full_hidden)) + pool_keys, pool_indices = self.indexer.prepare_pools(full_hidden) + if self.backend.attn == "cudnn": + return self._forward_document_cudnn( + full_hidden, + latent, + pool_keys, + pool_indices, + query_start, + query_end, + ) + + expanded = self.kv_b_proj(latent).view(1, length, self.num_heads, self.qk_head_dim + self.v_head_dim) + key, value = expanded.split([self.qk_head_dim, self.v_head_dim], dim=-1) + key, value = key.transpose(1, 2), value.transpose(1, 2) + chunks = [] + for start in range(query_start, query_end, self.query_chunk_size): + end = min(start + self.query_chunk_size, query_end) + query_hidden = full_hidden[:, start:end] + q_resid = self.q_a_layernorm(self.q_a_proj(query_hidden)) + query = self.q_b_proj(q_resid).view(1, end - start, self.num_heads, self.qk_head_dim).transpose(1, 2) + positions = torch.arange(start, end, device=full_hidden.device) + indices = self.indexer.select( + query_hidden, + q_resid, + positions, + pool_keys, + pool_indices, + length, + ) + valid = indices.ge(0) & indices.lt(length) + safe = indices.clamp(0, length - 1) + selected_counts = torch.zeros((1, end - start, length), dtype=torch.int32, device=full_hidden.device) + selected_counts.scatter_add_(-1, safe.long(), valid.to(torch.int32)) + selected = selected_counts.ne(0) + attn = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=selected.unsqueeze(1), + dropout_p=self.config.attention_dropout if self.training else 0.0, + scale=self.scaling, + ) + chunks.append(attn.transpose(1, 2).reshape(1, end - start, -1)) + return self.o_proj(torch.cat(chunks, dim=1)) + + def _forward_document_cudnn( + self, + full_hidden: torch.Tensor, + latent: torch.Tensor, + pool_keys: torch.Tensor, + pool_indices: torch.Tensor, + query_start: int, + query_end: int, + ) -> torch.Tensor: + """Run absorbed latent attention through FlashMLA/cuDNN for one document. + + Args: + full_hidden: Unpadded document states with shape ``[1, key_tokens, hidden]``. + latent: Normalized shared latent K/V with shape + ``[1, key_tokens, 512]``. + pool_keys: KPool-compressed index keys with shape + ``[complete_pools, index_head_dim]``. + pool_indices: Document-local token indices with shape + ``[complete_pools, index_kpool]``. + query_start: Inclusive document-local index of the first local query. + query_end: Exclusive document-local index of the final local query. + + Returns: + Projected attention output with shape + ``[1, query_end - query_start, hidden]``. + """ + if not is_cudnn_sparse_attention_available(): + raise RuntimeError( + "backend.attn='cudnn' requires the optional cuDNN sparse-attention " + "and FlashMLA runtimes, but they are unavailable in this environment." + ) + + weight = self.kv_b_proj.weight.view( + self.num_heads, + self.qk_head_dim + self.v_head_dim, + self.kv_lora_rank, + ) + w_kc, w_vc = weight.split([self.qk_head_dim, self.v_head_dim], dim=1) + absorbed_queries = [] + selected_indices = [] + length = full_hidden.shape[1] + for start in range(query_start, query_end, self.query_chunk_size): + end = min(start + self.query_chunk_size, query_end) + query_hidden = full_hidden[:, start:end] + q_resid = self.q_a_layernorm(self.q_a_proj(query_hidden)) + query = self.q_b_proj(q_resid).view(1, end - start, self.num_heads, self.qk_head_dim) + absorbed_queries.append(torch.einsum("bqhd,hdc->bqhc", query, w_kc.to(query.dtype)).squeeze(0)) + positions = torch.arange(start, end, device=full_hidden.device) + indices = self.indexer.select( + query_hidden, + q_resid, + positions, + pool_keys, + pool_indices, + length, + ) + selected_indices.append(indices.squeeze(0).unsqueeze(1)) + + latent_output = cudnn_sparse_attention( + torch.cat(absorbed_queries, dim=0).contiguous(), + latent.squeeze(0).unsqueeze(1).contiguous(), + torch.cat(selected_indices, dim=0).contiguous(), + self.scaling, + all_rows_nonempty=self.indexer.always_select_tail, + ) + attention_output = torch.einsum("qhc,hvc->qhv", latent_output, w_vc.to(latent_output.dtype)) + return self.o_proj(attention_output.reshape(1, query_end - query_start, -1)) + + @torch.no_grad() + def init_weights(self, buffer_device: torch.device, init_std: float) -> None: + """Initialize MLA and indexer parameters.""" + with buffer_device: + for module in (self.q_a_proj, self.q_b_proj, self.kv_a_proj_with_mqa, self.kv_b_proj, self.o_proj): + nn.init.normal_(module.weight, mean=0.0, std=init_std) + self.q_a_layernorm.reset_parameters() + self.kv_a_layernorm.reset_parameters() + self.indexer.init_weights(buffer_device, init_std) + + +class Glm5NextDecoderLayer(nn.Module): + """One mHC decoder block with KDA/DSA and dense/MoE feed-forward.""" + + def __init__( + self, + config: Glm5NextTextConfig, + layer_idx: int, + moe_config: MoEConfig, + backend: BackendConfig, + ) -> None: + super().__init__() + self.layer_idx = layer_idx + self.block_type = config.layer_types[layer_idx] + self.is_linear_attn = self.block_type == "linear_attention" + self.is_moe_layer = config.mlp_layer_types[layer_idx] == "sparse" + self.self_attn = ( + Glm5NextLinearAttention(config, layer_idx) + if self.is_linear_attn + else Glm5NextSparseAttention(config, layer_idx, backend) + ) + dtype = get_dtype(getattr(config, "torch_dtype", None), torch.bfloat16) + self.mlp = ( + MoE(moe_config, backend) + if self.is_moe_layer + else MLP( + config.hidden_size, + config.intermediate_size, + backend.linear, + dtype=dtype, + swiglu_limit=config.swiglu_limit, + ) + ) + self.input_layernorm = Glm5NextRMSNorm(config.hidden_size, config.rms_norm_eps, dtype) + self.post_attention_layernorm = Glm5NextRMSNorm(config.hidden_size, config.rms_norm_eps, dtype) + self.attn_hc = Glm5NextHyperConnection(config) + self.ffn_hc = Glm5NextHyperConnection(config) + + def forward( + self, + hidden_streams: torch.Tensor, + *, + packed_context: Glm5NextPackedContext, + padding_mask: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor: + """Transform residual streams ``[batch, local_sequence, hc_mult, hidden]``.""" + dtype = hidden_streams.dtype + residual = hidden_streams + post, comb, collapsed = self.attn_hc(hidden_streams) + update = self.self_attn( + self.input_layernorm(collapsed), + packed_context=packed_context, + padding_mask=padding_mask, + **kwargs, + ) + hidden_streams = post.to(dtype).unsqueeze(-1) * update.unsqueeze(-2) + torch.matmul( + comb.to(dtype).transpose(-1, -2), residual + ) + residual = hidden_streams + post, comb, collapsed = self.ffn_hc(hidden_streams) + update = self.post_attention_layernorm(collapsed) + update = self.mlp(update, padding_mask) if self.is_moe_layer else self.mlp(update) + return post.to(dtype).unsqueeze(-1) * update.unsqueeze(-2) + torch.matmul( + comb.to(dtype).transpose(-1, -2), residual + ) + + def update_moe_gate_bias(self) -> None: + """Update the correction bias for this layer's learned MoE router.""" + if self.is_moe_layer: + moe = unwrap_checkpoint_wrapper(self.mlp) + if isinstance(moe, MoE) and moe.gate.bias_update_factor > 0: + moe.gate.update_bias() + + @torch.no_grad() + def init_weights(self, buffer_device: torch.device, init_std: float) -> None: + """Initialize all decoder children.""" + self.input_layernorm.reset_parameters() + self.post_attention_layernorm.reset_parameters() + self.attn_hc.init_weights(buffer_device, init_std) + self.ffn_hc.init_weights(buffer_device, init_std) + self.self_attn.init_weights(buffer_device, init_std) + self.mlp.init_weights(buffer_device, init_std) + + +__all__ = [ + "Glm5NextDecoderLayer", + "Glm5NextHyperConnection", + "Glm5NextKPoolIndexer", + "Glm5NextLinearAttention", + "Glm5NextRMSNorm", + "Glm5NextSparseAttention", +] diff --git a/nemo_automodel/components/models/glm5_next/model.py b/nemo_automodel/components/models/glm5_next/model.py new file mode 100644 index 0000000000..c419ddff74 --- /dev/null +++ b/nemo_automodel/components/models/glm5_next/model.py @@ -0,0 +1,461 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native AutoModel implementation of GLM-5.3-Flash.""" + +from __future__ import annotations + +import copy +from dataclasses import dataclass +from functools import partial +from typing import Any + +import torch +from torch import nn +from transformers.modeling_outputs import CausalLMOutputWithPast + +from nemo_automodel.components.distributed.context_parallel.sharder import ( + ContextParallelSharder, + contiguous_local_indices, +) +from nemo_automodel.components.models.common import BackendConfig, initialize_linear_module +from nemo_automodel.components.models.common.hf_checkpointing_mixin import HFCheckpointingMixin +from nemo_automodel.components.models.common.tie_word_embeddings import ( + TieSupport, + reject_unsupported_tie_word_embeddings, +) +from nemo_automodel.components.models.common.utils import cast_model_to_dtype, compute_lm_head_logits +from nemo_automodel.components.models.glm5_next.config import Glm5NextConfig, Glm5NextTextConfig +from nemo_automodel.components.models.glm5_next.cp import ( + Glm5NextPackedContext, + doc_ids_from_cu_seqlens, + shard_batch_for_glm5_next_cp, +) +from nemo_automodel.components.models.glm5_next.layers import ( + Glm5NextDecoderLayer, + Glm5NextRMSNorm, +) +from nemo_automodel.components.models.glm5_next.vision import Glm5NextVisionModel, Glm5NextVisionOutput +from nemo_automodel.components.moe.config import MoEConfig +from nemo_automodel.components.moe.fsdp_mixin import MoEFSDPSyncMixin +from nemo_automodel.shared.utils import dtype_from_str as get_dtype + + +def build_glm5_next_moe_config( + config: Glm5NextTextConfig, + dtype: torch.dtype, + overrides: dict[str, Any] | None = None, +) -> MoEConfig: + """Translate the GLM router/expert contract to AutoModel's grouped MoE.""" + values = dict( + dim=config.hidden_size, + inter_dim=config.intermediate_size, + moe_inter_dim=config.moe_intermediate_size, + n_routed_experts=config.n_routed_experts, + n_shared_experts=config.n_shared_experts, + n_activated_experts=config.num_experts_per_tok, + n_expert_groups=config.n_group, + n_limited_groups=config.topk_group, + train_gate=True, + gate_bias_update_factor=1e-3, + # HF uses the correction bias only to select experts; routing weights + # are gathered from the unbiased sigmoid scores. + score_func="sigmoid_with_bias", + route_scale=config.routed_scaling_factor, + aux_loss_coeff=0.0, + norm_topk_prob=config.norm_topk_prob, + router_bias=False, + expert_bias=False, + expert_activation="swiglu", + apply_router_weight_after_down=True, + swiglu_limit=config.swiglu_limit, + softmax_before_topk=False, + router_weights_fp32=True, + router_weight_uses_score_correction_bias=False, + shared_expert_gate=False, + shared_expert_inter_dim=config.moe_intermediate_size, + force_e_score_correction_bias=True, + dtype=dtype, + ) + if overrides: + values.update(overrides) + return MoEConfig(**values) + + +def _packed_context_from_inputs( + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None, + padding_mask: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, + packed_seq_ids: torch.Tensor | None, +) -> Glm5NextPackedContext: + """Build one global document map for a non-CP forward.""" + sequence = input_ids.shape[1] + if packed_seq_ids is not None: + doc_ids = packed_seq_ids.to(torch.int32) + if doc_ids.ndim == 1: + doc_ids = doc_ids.unsqueeze(0) + elif cu_seqlens is not None: + doc_ids = doc_ids_from_cu_seqlens(cu_seqlens, sequence) + elif attention_mask is not None and attention_mask.ndim == 2: + doc_ids = attention_mask.to(torch.int32) + else: + doc_ids = torch.ones(input_ids.shape[:2], dtype=torch.int32, device=input_ids.device) + if padding_mask is not None: + doc_ids.masked_fill_(padding_mask.bool(), 0) + return Glm5NextPackedContext(doc_ids=doc_ids, original_seq_len=sequence) + + +class Glm5NextTextModel(nn.Module): + """Embedding, mHC decoder stack, mean stream collapse and final RMSNorm.""" + + def __init__( + self, + config: Glm5NextTextConfig, + backend: BackendConfig, + *, + moe_config: MoEConfig | None = None, + moe_overrides: dict[str, Any] | None = None, + ) -> None: + super().__init__() + if moe_config is not None and moe_overrides is not None: + raise ValueError("Pass either moe_config or moe_overrides, not both") + self.config = config + self.backend = backend + dtype = get_dtype(getattr(config, "torch_dtype", getattr(config, "dtype", None)), torch.bfloat16) + self.moe_config = moe_config or build_glm5_next_moe_config(config, dtype, moe_overrides) + self.padding_idx = config.pad_token_id + self.embed_tokens = nn.Embedding( + config.vocab_size, + config.hidden_size, + padding_idx=config.pad_token_id, + dtype=dtype, + ) + self.layers = nn.ModuleDict( + { + str(layer_idx): Glm5NextDecoderLayer(config, layer_idx, self.moe_config, backend) + for layer_idx in range(config.num_hidden_layers) + } + ) + self.norm = Glm5NextRMSNorm(config.hidden_size, config.rms_norm_eps, dtype) + + def forward( + self, + input_ids: torch.Tensor | None = None, + *, + inputs_embeds: torch.Tensor | None = None, + glm5_next_packed_context: Glm5NextPackedContext, + padding_mask: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor: + """Run ``[batch, local_sequence]`` ids/embeddings through the text model.""" + if (input_ids is None) == (inputs_embeds is None): + raise ValueError("Specify exactly one of input_ids and inputs_embeds") + hidden = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds + hidden = hidden.unsqueeze(2).expand(-1, -1, self.config.hc_mult, -1).contiguous() + for layer in self.layers.values(): + hidden = layer( + hidden, + packed_context=glm5_next_packed_context, + padding_mask=padding_mask, + **kwargs, + ) + return self.norm(hidden.mean(dim=2)) + + def update_moe_gate_bias(self) -> None: + """Update every sparse layer's no-aux-loss routing correction bias.""" + for layer in self.layers.values(): + layer.update_moe_gate_bias() + + @torch.no_grad() + def init_weights(self, buffer_device: torch.device) -> None: + """Initialize a checkpoint-free text model on ``buffer_device``.""" + init_std = self.config.initializer_range + with buffer_device: + nn.init.normal_(self.embed_tokens.weight, mean=0.0, std=init_std) + if self.padding_idx is not None: + self.embed_tokens.weight[self.padding_idx].zero_() + self.norm.reset_parameters() + for layer in self.layers.values(): + layer.init_weights(buffer_device, init_std) + + +class Glm5NextModel(nn.Module): + """Checkpoint-layout container for ``visual`` and ``language_model``.""" + + def __init__( + self, + config: Glm5NextConfig, + backend: BackendConfig, + *, + moe_config: MoEConfig | None = None, + moe_overrides: dict[str, Any] | None = None, + ) -> None: + super().__init__() + self.visual = Glm5NextVisionModel(config.vision_config) + self.language_model = Glm5NextTextModel( + config.text_config, + backend, + moe_config=moe_config, + moe_overrides=moe_overrides, + ) + + def get_image_features(self, pixel_values: torch.Tensor, image_grid_thw: torch.Tensor) -> Glm5NextVisionOutput: + """Encode image patches and split-free concatenated features.""" + return self.visual(pixel_values.to(self.visual.dtype), image_grid_thw) + + +class Glm5NextForConditionalGeneration(HFCheckpointingMixin, nn.Module, MoEFSDPSyncMixin): + """Trainable GLM-5.3 VLM with EP and contiguous packed CP support.""" + + tie_word_embeddings_support: TieSupport = TieSupport.UNTIED_ONLY + _skip_init_weights_on_load = True + _owns_cp_attention = True + _owns_packed_attention = True + _packed_cp_attn_backends = ("sdpa", "cudnn") + _keep_in_fp32_modules_strict = [ + "_fp32_params", + "e_score_correction_bias", + "rotary_pos_emb", + ] + cp_mesh = None + + @dataclass(frozen=True) + class ModelCapabilities: + """Parallel axes intentionally supported for the released checkpoint.""" + + supports_tp: bool = False + supports_cp: bool = True + supports_pp: bool = False + supports_ep: bool = True + supports_thd: bool = True + + @classmethod + def from_config( + cls, + config: Glm5NextConfig, + moe_config: MoEConfig | None = None, + backend: BackendConfig | None = None, + **kwargs: Any, + ) -> "Glm5NextForConditionalGeneration": + """Construct from an already resolved native config.""" + return cls(config, moe_config=moe_config, backend=backend, **kwargs) + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: str, + *model_args: Any, + **kwargs: Any, + ) -> "Glm5NextForConditionalGeneration": + """Resolve the local config; checkpoint loading is owned by AutoModel.""" + config = Glm5NextConfig.from_pretrained(pretrained_model_name_or_path) + return cls.from_config(config, *model_args, **kwargs) + + def __init__( + self, + config: Glm5NextConfig, + moe_config: MoEConfig | None = None, + backend: BackendConfig | None = None, + **kwargs: Any, + ) -> None: + super().__init__() + reject_unsupported_tie_word_embeddings(type(self), config) + self.config = config + self.backend = copy.copy(backend) if backend is not None else BackendConfig() + if self.backend.gate_precision is None: + self.backend.gate_precision = torch.float32 + moe_overrides = kwargs.pop("moe_overrides", None) + self.model = Glm5NextModel( + config, + self.backend, + moe_config=moe_config, + moe_overrides=moe_overrides, + ) + text_config = config.text_config + dtype = get_dtype( + getattr(text_config, "torch_dtype", getattr(text_config, "dtype", None)), + torch.bfloat16, + ) + self.lm_head = initialize_linear_module( + self.backend.linear, + text_config.hidden_size, + text_config.vocab_size, + bias=False, + dtype=dtype, + ) + self.vocab_size = text_config.vocab_size + if self.backend.enable_hf_state_dict_adapter: + from nemo_automodel.components.models.glm5_next.state_dict_adapter import Glm5NextStateDictAdapter + + self.state_dict_adapter = Glm5NextStateDictAdapter( + config, + self.model.language_model.moe_config, + self.backend, + dtype=dtype, + ) + + @property + def language_model(self) -> Glm5NextTextModel: + """Expose the text module through the multimodal discovery protocol.""" + return self.model.language_model + + def get_input_embeddings(self) -> nn.Module: + """Return the token embedding table.""" + return self.model.language_model.embed_tokens + + def set_input_embeddings(self, value: nn.Module) -> None: + """Replace the token embedding table.""" + self.model.language_model.embed_tokens = value + + def get_output_embeddings(self) -> nn.Module: + """Return the untied language-model head.""" + return self.lm_head + + def set_output_embeddings(self, value: nn.Module) -> None: + """Replace the language-model head.""" + self.lm_head = value + + def get_image_features(self, pixel_values: torch.Tensor, image_grid_thw: torch.Tensor) -> Glm5NextVisionOutput: + """Return raw and merged features for flattened image patches.""" + return self.model.get_image_features(pixel_values, image_grid_thw) + + def _embed_and_splice( + self, + input_ids: torch.Tensor, + pixel_values: torch.Tensor | None, + image_grid_thw: torch.Tensor | None, + ) -> torch.Tensor: + """Embed the full sequence and replace image placeholder positions.""" + embeddings = self.get_input_embeddings()(input_ids) + if pixel_values is None: + return embeddings + if image_grid_thw is None: + raise ValueError("image_grid_thw is required when pixel_values are provided") + features = self.get_image_features(pixel_values, image_grid_thw).pooler_output + mask = input_ids == self.config.image_token_id + expected = int(mask.sum().item()) + if features.shape[0] != expected: + raise ValueError(f"GLM-5.3 produced {features.shape[0]} image tokens for {expected} placeholders") + embeddings = embeddings.clone() + embeddings[mask] = features.to(device=embeddings.device, dtype=embeddings.dtype) + return embeddings + + def forward( + self, + input_ids: torch.Tensor | None = None, + *, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + padding_mask: torch.Tensor | None = None, + pixel_values: torch.Tensor | None = None, + image_grid_thw: torch.Tensor | None = None, + pixel_values_videos: torch.Tensor | None = None, + logits_to_keep: int | torch.Tensor = 0, + output_hidden_states: bool | None = None, + **kwargs: Any, + ) -> CausalLMOutputWithPast: + """Run image splice, contiguous CP slicing, text decoding and lm head.""" + del position_ids + if pixel_values_videos is not None: + raise NotImplementedError("GLM-5.3 AutoModel onboarding currently supports images, not video training") + is_thd = kwargs.get("qkv_format") == "thd" + if input_ids is not None and input_ids.ndim == 1: + input_ids = input_ids.unsqueeze(0) + if inputs_embeds is not None and inputs_embeds.ndim == 2: + inputs_embeds = inputs_embeds.unsqueeze(0) + if padding_mask is not None and padding_mask.ndim == 1: + padding_mask = padding_mask.unsqueeze(0) + if input_ids is None and inputs_embeds is None: + raise ValueError("input_ids or inputs_embeds is required") + + context = kwargs.pop("glm5_next_packed_context", None) + context_ids = input_ids + if context_ids is None: + context_ids = torch.zeros(inputs_embeds.shape[:2], dtype=torch.long, device=inputs_embeds.device) + if context is None: + context = _packed_context_from_inputs( + context_ids, + attention_mask, + padding_mask, + kwargs.get("cu_seqlens"), + kwargs.get("_packed_seq_ids"), + ) + + if inputs_embeds is None: + inputs_embeds = self._embed_and_splice(input_ids, pixel_values, image_grid_thw) + elif pixel_values is not None: + raise ValueError("pixel_values cannot be combined with precomputed inputs_embeds") + + padded_length = context.doc_ids.shape[1] + if inputs_embeds.shape[1] < padded_length: + pad = inputs_embeds.new_zeros( + inputs_embeds.shape[0], padded_length - inputs_embeds.shape[1], inputs_embeds.shape[2] + ) + inputs_embeds = torch.cat((inputs_embeds, pad), dim=1) + local_length = context.local_seq_len + if inputs_embeds.shape[1] != local_length or context.seq_start: + inputs_embeds = inputs_embeds[:, context.seq_start : context.seq_start + local_length].contiguous() + if padding_mask is None: + padding_mask = context.local_doc_ids <= 0 + + hidden = self.model.language_model( + inputs_embeds=inputs_embeds, + glm5_next_packed_context=context, + padding_mask=padding_mask, + **kwargs, + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else bool(getattr(self.config.text_config, "output_hidden_states", False)) + ) + return compute_lm_head_logits( + self.lm_head, + hidden, + logits_to_keep, + is_thd=is_thd, + output_hidden_states=output_hidden_states, + ) + + def prepare_model_inputs_for_cp(self, batch: dict[str, Any], *, num_chunks: int = 1) -> dict[str, Any]: + """Install GLM's contiguous packed sharder while leaving media and ids global.""" + del batch, num_chunks + return { + "cp_sharder": ContextParallelSharder( + shard_batch=partial(shard_batch_for_glm5_next_cp, shard_primary=False), + local_token_global_indices=contiguous_local_indices, + ) + } + + def update_moe_gate_bias(self) -> None: + """Update no-aux-loss router correction biases after an optimizer step.""" + self.model.language_model.update_moe_gate_bias() + + @torch.no_grad() + def initialize_weights( + self, + buffer_device: torch.device | None = None, + dtype: torch.dtype = torch.bfloat16, + ) -> None: + """Initialize all tensors for checkpoint-free construction.""" + buffer_device = buffer_device or torch.device( + f"cuda:{torch.cuda.current_device()}" if torch.cuda.is_available() else "cpu" + ) + self.model.language_model.init_weights(buffer_device) + self.model.visual.init_weights(buffer_device, self.config.text_config.initializer_range) + final_std = self.config.text_config.hidden_size**-0.5 + with buffer_device: + nn.init.trunc_normal_(self.lm_head.weight, mean=0.0, std=final_std, a=-3 * final_std, b=3 * final_std) + cast_model_to_dtype(self, dtype, skip_modules=("_fp32_params",)) + + +ModelClass = Glm5NextForConditionalGeneration + +__all__ = [ + "Glm5NextForConditionalGeneration", + "Glm5NextModel", + "Glm5NextTextModel", + "build_glm5_next_moe_config", +] diff --git a/nemo_automodel/components/models/glm5_next/processing.py b/nemo_automodel/components/models/glm5_next/processing.py new file mode 100644 index 0000000000..99f2ad0486 --- /dev/null +++ b/nemo_automodel/components/models/glm5_next/processing.py @@ -0,0 +1,103 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Build the GLM-5.3 image processor on Transformers versions before 5.16.""" + +from __future__ import annotations + +import json +import os +from typing import Any + +from huggingface_hub.utils import EntryNotFoundError +from transformers import AutoTokenizer +from transformers.models.glm46v.processing_glm46v import Glm46VProcessor +from transformers.processing_utils import ProcessorMixin + +from nemo_automodel.components.models.glm5_next.image_processing import Glm5NextImageProcessor +from nemo_automodel.shared.import_utils import safe_import_from + +_, Glm46VVideoProcessor = safe_import_from( + "transformers.models.glm46v.video_processing_glm46v", + "Glm46VVideoProcessor", + msg="GLM-5.3 processor construction requires torchvision. Install the `vlm` extra.", +) + +_MEDIA_REMINDER = ( + '{{- "You are unable to process this " ~ media_type ~ ' + '" because you don\'t have multi-modal input ability. Try different methods." }}' +) +_IMAGE_PLACEHOLDER = """{%- if media_type == 'image' -%} + {{- "<|begin_of_image|><|image|><|end_of_image|>" }} + {%- else -%} + {{- "You are unable to process this " ~ media_type ~ " because you don't have multi-modal input ability. Try different methods." }} + {%- endif -%}""" + + +def _load_processor_config(path_or_id: str, **kwargs: Any) -> dict[str, Any]: + local = os.path.join(path_or_id, "processor_config.json") + if not os.path.isfile(local): + from huggingface_hub import hf_hub_download + + hub_kwargs = {key: kwargs[key] for key in ("cache_dir", "revision", "token") if key in kwargs} + local = hf_hub_download(path_or_id, "processor_config.json", **hub_kwargs) + with open(local, encoding="utf-8") as stream: + return json.load(stream) + + +def _load_chat_template(path_or_id: str, **kwargs: Any) -> str | None: + local = os.path.join(path_or_id, "chat_template.jinja") + if not os.path.isfile(local): + if os.path.isdir(path_or_id): + return None + try: + from huggingface_hub import hf_hub_download + + hub_kwargs = {key: kwargs[key] for key in ("cache_dir", "revision", "token") if key in kwargs} + local = hf_hub_download(path_or_id, "chat_template.jinja", **hub_kwargs) + except EntryNotFoundError: + return None + with open(local, encoding="utf-8") as stream: + return stream.read() + + +def _enable_image_placeholders(template: str | None) -> str | None: + """Render image content as GLM image tokens while preserving the shipped template. + + The initial GLM-5.3-Flash checkpoint template renders every media block as a + text-only capability reminder. That leaves ``Glm46VProcessor.__call__`` no + ``<|image|>`` token to expand even though it receives and patchifies the image. + MedPix fine-tuning requires the native begin/image/end token triplet; video + retains the checkpoint's reminder because this onboarding is image-only. + """ + if template is None: + return None + if "<|image|>" in template: + return template + if _MEDIA_REMINDER not in template: + raise ValueError("GLM-5.3 chat template has no recognized media rendering branch") + return template.replace(_MEDIA_REMINDER, _IMAGE_PLACEHOLDER, 1) + + +def build_glm5_next_processor(pretrained_model_name_or_path: str, **kwargs: Any) -> ProcessorMixin: + """Create the image-only GLM-5.3 processor used by MedPix recipes.""" + processor_config = _load_processor_config(pretrained_model_name_or_path, **kwargs) + image_kwargs = dict(processor_config.get("image_processor", {})) + image_kwargs.pop("image_processor_type", None) + image_processor = Glm5NextImageProcessor(**image_kwargs) + video_kwargs = dict(processor_config.get("video_processor", {})) + video_kwargs.pop("video_processor_type", None) + video_processor = Glm46VVideoProcessor(**video_kwargs) + tokenizer_kwargs = dict(kwargs) + tokenizer_kwargs.pop("trust_remote_code", None) + tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, **tokenizer_kwargs) + template = _enable_image_placeholders(_load_chat_template(pretrained_model_name_or_path, **kwargs)) + return Glm46VProcessor( + image_processor=image_processor, + tokenizer=tokenizer, + video_processor=video_processor, + chat_template=template, + ) + + +__all__ = ["build_glm5_next_processor"] diff --git a/nemo_automodel/components/models/glm5_next/state_dict_adapter.py b/nemo_automodel/components/models/glm5_next/state_dict_adapter.py new file mode 100644 index 0000000000..e19aff3a79 --- /dev/null +++ b/nemo_automodel/components/models/glm5_next/state_dict_adapter.py @@ -0,0 +1,287 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Checkpoint conversion for the released GLM-5.3-Flash VLM. + +The checkpoint already uses the native vision/text prefix layout. Conversion is +needed for grouped EP experts, flat mHC/KDA parameters, the extra MTP layer, and +128x128 block-scaled FP8 training weights. +""" + +from __future__ import annotations + +import re +from typing import Any + +import torch +from torch.distributed.device_mesh import DeviceMesh + +from nemo_automodel.components.checkpoint.state_dict_adapter import StateDictAdapter +from nemo_automodel.components.models.common import BackendConfig +from nemo_automodel.components.models.glm5_next.config import Glm5NextConfig +from nemo_automodel.components.moe.config import MoEConfig +from nemo_automodel.components.moe.state_dict_mixin import MoESplitExpertsStateDictMixin +from nemo_automodel.components.moe.state_dict_utils import is_dtensor + +_BLOCK_SIZE = 128 +_FP8_WEIGHT = re.compile( + r"^model\.language_model\.layers\.\d+\.(?:" + r"self_attn\.(?:q_a_proj|q_b_proj|kv_a_proj_with_mqa)" + r"|mlp\.(?:gate|up|down)_proj" + r"|mlp\.experts\.\d+\.(?:gate|up|down)_proj" + r"|mlp\.shared_experts\.(?:gate|up|down)_proj" + r")\.weight$" +) +_SPARSE_O_WEIGHT = re.compile(r"^model\.language_model\.layers\.(\d+)\.self_attn\.o_proj\.weight$") +_HC_KEY = re.compile(r"^(model\.language_model\.layers\.\d+)\.hc_(attn|ffn)_(fn|base|scale)$") +_NATIVE_HC_KEY = re.compile( + r"^(model\.language_model\.layers\.\d+)\.(attn_hc|ffn_hc)(?:\._fp32_params)?\.(fn|base|scale)$" +) +_KDA_PARAMETER = re.compile(r"^(model\.language_model\.layers\.\d+\.self_attn)\.(A_log|dt_bias)$") +_NATIVE_KDA_PARAMETER = re.compile(r"^(model\.language_model\.layers\.\d+\.self_attn)\._fp32_params\.(A_log|dt_bias)$") + + +def _scale_shape(weight: torch.Tensor) -> tuple[int, int]: + return ( + (weight.shape[-2] + _BLOCK_SIZE - 1) // _BLOCK_SIZE, + (weight.shape[-1] + _BLOCK_SIZE - 1) // _BLOCK_SIZE, + ) + + +def _scale_placeholder(weight: torch.Tensor) -> torch.Tensor: + """Create the global FP8 block-scale load destination for a 2-D weight.""" + local = weight.to_local() if is_dtensor(weight) else weight + return torch.ones(_scale_shape(weight), dtype=torch.float32, device=local.device) + + +def _local_shard_offsets(tensor: torch.Tensor) -> tuple[int, ...]: + """Return the global start coordinate of a DTensor's contiguous local shard.""" + from torch.distributed.tensor import Shard + + offsets = [0] * tensor.ndim + current_shape = list(tensor.shape) + for mesh_dim, placement in enumerate(tensor.placements): + if not isinstance(placement, Shard) or placement.dim >= tensor.ndim: + continue + shard_dim = placement.dim + local_size, relative_offset = Shard.local_shard_size_and_offset( + current_shape[shard_dim], + tensor.device_mesh.size(mesh_dim), + tensor.device_mesh.get_local_rank(mesh_dim=mesh_dim), + ) + offsets[shard_dim] += int(relative_offset) + current_shape[shard_dim] = int(local_size) + return tuple(offsets) + + +def _apply_local_block_scales( + local_weight: torch.Tensor, + local_scale: torch.Tensor, + local_offsets: tuple[int, int], + dtype: torch.dtype, +) -> torch.Tensor: + """Apply global-grid block scales to one possibly misaligned local shard.""" + rows, cols = local_weight.shape + row_offset = local_offsets[0] % _BLOCK_SIZE + col_offset = local_offsets[1] % _BLOCK_SIZE + expected = ( + (row_offset + rows + _BLOCK_SIZE - 1) // _BLOCK_SIZE, + (col_offset + cols + _BLOCK_SIZE - 1) // _BLOCK_SIZE, + ) + if tuple(local_scale.shape) != expected: + raise ValueError( + f"FP8 scale shape {tuple(local_scale.shape)} does not cover local weight " + f"{tuple(local_weight.shape)} at global offset {local_offsets} (expected {expected})" + ) + expanded = local_scale.float().repeat_interleave(_BLOCK_SIZE, 0).repeat_interleave(_BLOCK_SIZE, 1) + scale = expanded[row_offset : row_offset + rows, col_offset : col_offset + cols] + return (local_weight.float() * scale).to(dtype) + + +def dequantize_block_fp8( + weight: torch.Tensor, + scale_inv: torch.Tensor, + *, + dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Dequantize an e4m3 weight with 128x128 fp32 inverse scales.""" + weight_is_dtensor = is_dtensor(weight) + scale_is_dtensor = is_dtensor(scale_inv) + local_weight = weight.to_local() if weight_is_dtensor else weight + local_scale = scale_inv.to_local() if scale_is_dtensor else scale_inv + local_offsets = _local_shard_offsets(weight) if weight_is_dtensor else (0, 0) + block_starts = tuple(offset // _BLOCK_SIZE for offset in local_offsets) + block_ends = tuple( + (offset + size + _BLOCK_SIZE - 1) // _BLOCK_SIZE for offset, size in zip(local_offsets, local_weight.shape) + ) + expected_local_scale = tuple(end - start for start, end in zip(block_starts, block_ends)) + + if scale_is_dtensor: + scale_offsets = _local_shard_offsets(scale_inv) + if scale_offsets != block_starts or tuple(local_scale.shape) != expected_local_scale: + raise ValueError( + "FP8 scale DTensor shard does not cover the corresponding weight shard: " + f"scale offset/shape={scale_offsets}/{tuple(local_scale.shape)}, " + f"expected={block_starts}/{expected_local_scale}" + ) + elif tuple(scale_inv.shape) == _scale_shape(weight): + local_scale = scale_inv[ + block_starts[0] : block_ends[0], + block_starts[1] : block_ends[1], + ] + elif tuple(local_scale.shape) != expected_local_scale: + raise ValueError( + f"FP8 scale shape {tuple(local_scale.shape)} does not match global weight " + f"{tuple(weight.shape)} or its local block coverage {expected_local_scale}" + ) + + output = _apply_local_block_scales( + local_weight, + local_scale.to(local_weight.device), + local_offsets, + dtype, + ) + if weight_is_dtensor: + from torch.distributed.tensor import DTensor + + return DTensor.from_local( + output, + weight.device_mesh, + weight.placements, + shape=weight.shape, + stride=weight.stride(), + ) + return output + + +def _hf_to_native_key(key: str) -> str: + match = _HC_KEY.match(key) + if match: + site = "attn_hc" if match.group(2) == "attn" else "ffn_hc" + holder = "._fp32_params" if match.group(3) in ("base", "scale") else "" + return f"{match.group(1)}.{site}{holder}.{match.group(3)}" + match = _KDA_PARAMETER.match(key) + if match: + return f"{match.group(1)}._fp32_params.{match.group(2)}" + return key + + +def _native_to_hf_key(key: str) -> str: + match = _NATIVE_HC_KEY.match(key) + if match: + site = "attn" if match.group(2) == "attn_hc" else "ffn" + return f"{match.group(1)}.hc_{site}_{match.group(3)}" + match = _NATIVE_KDA_PARAMETER.match(key) + if match: + return f"{match.group(1)}.{match.group(2)}" + return key + + +class Glm5NextStateDictAdapter(MoESplitExpertsStateDictMixin, StateDictAdapter): + """Convert GLM split experts and FP8 weights to trainable grouped BF16.""" + + def __init__( + self, + config: Glm5NextConfig, + moe_config: MoEConfig, + backend: BackendConfig, + dtype: torch.dtype = torch.bfloat16, + ) -> None: + self.config = config + self.moe_config = moe_config + self.backend = backend + self.dtype = dtype + self._uses_model_prefix = True + + @property + def _hf_prefix(self) -> str: + return "model.language_model." + + @property + def _expert_path_segment(self) -> str: + return "mlp.experts" + + def _dequantize(self, state_dict: dict[str, Any]) -> None: + scale_keys = [] + for key, value in list(state_dict.items()): + scale_key = key + "_scale_inv" + if key.endswith(".weight") and scale_key in state_dict: + state_dict[key] = dequantize_block_fp8(value, state_dict[scale_key], dtype=self.dtype) + scale_keys.append(scale_key) + for key in scale_keys: + state_dict.pop(key, None) + + def _is_fp8_weight(self, key: str) -> bool: + """Match the checkpoint's quantized matrices, including DSA-only output projections.""" + if _FP8_WEIGHT.match(key): + return True + match = _SPARSE_O_WEIGHT.match(key) + if match is None: + return False + return self.config.text_config.layer_types[int(match.group(1))] != "linear_attention" + + def from_hf( + self, + hf_state_dict: dict[str, Any], + device_mesh: DeviceMesh | None = None, + **kwargs: Any, + ) -> dict[str, Any]: + """Dequantize, drop MTP, route flat parameters and aggregate experts.""" + del kwargs + layer_limit = self.config.text_config.num_hidden_layers + mtp_prefix = f"model.language_model.layers.{layer_limit}." + for key in list(hf_state_dict): + if key.startswith(mtp_prefix): + hf_state_dict.pop(key) + self._dequantize(hf_state_dict) + for key in list(hf_state_dict): + value = hf_state_dict.pop(key) + native_key = _hf_to_native_key(key) + if native_key.endswith("._fp32_params.A_log"): + value = value.float() + elif native_key.endswith("._fp32_params.dt_bias") or native_key.endswith(".e_score_correction_bias"): + value = value.float() + hf_state_dict[native_key] = value + return self._from_hf_w_merged_experts(hf_state_dict, device_mesh) + + def to_hf( + self, + state_dict: dict[str, Any], + exclude_key_regex: str | None = None, + quantization: bool = False, + **kwargs: Any, + ) -> dict[str, Any]: + """Expand grouped experts and restore released checkpoint names.""" + output: dict[str, Any] = {} + for key, value in state_dict.items(): + for hf_key, hf_value in self.convert_single_tensor_to_hf( + key, + value, + exclude_key_regex=exclude_key_regex, + quantization=quantization, + **kwargs, + ): + output[hf_key] = hf_value + return output + + def convert_single_tensor_to_hf(self, fqn: str, tensor: Any, **kwargs: Any) -> list[tuple[str, Any]]: + """Convert one native tensor, including split expert and FP8 load targets.""" + exclude = kwargs.get("exclude_key_regex") + quantization = kwargs.get("quantization", False) + expert = self._convert_single_merged_expert_to_hf_split_experts(fqn, tensor, **kwargs) + result = expert if expert is not None else [(fqn, tensor)] + converted: list[tuple[str, Any]] = [] + for key, value in result: + key = _native_to_hf_key(key) + if exclude and re.match(exclude, key): + continue + if quantization and self._is_fp8_weight(key): + fp8_value = value.to(torch.float8_e4m3fn) + converted.append((key, fp8_value)) + converted.append((key + "_scale_inv", _scale_placeholder(value))) + else: + converted.append((key, value)) + return converted + + +__all__ = ["Glm5NextStateDictAdapter", "dequantize_block_fp8"] diff --git a/nemo_automodel/components/models/glm5_next/vision.py b/nemo_automodel/components/models/glm5_next/vision.py new file mode 100644 index 0000000000..0f7dbe8825 --- /dev/null +++ b/nemo_automodel/components/models/glm5_next/vision.py @@ -0,0 +1,268 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""GLM-5.3 image encoder with checkpoint-compatible module names.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.nn.functional as F +from torch import nn + +from nemo_automodel.components.models.glm5_next.config import Glm5NextVisionConfig +from nemo_automodel.components.models.glm5_next.layers import Glm5NextRMSNorm +from nemo_automodel.shared.utils import dtype_from_str as get_dtype + + +def _rotate_half(hidden_states: torch.Tensor) -> torch.Tensor: + first, second = hidden_states.chunk(2, dim=-1) + return torch.cat((-second, first), dim=-1) + + +def _vision_position_ids(grid_thw: torch.Tensor, merge_size: int) -> torch.Tensor: + """Return block-major H/W positions ``[vision_tokens, 2]``.""" + positions = [] + for temporal, height, width in grid_thw.tolist(): + hpos, wpos = torch.meshgrid( + torch.arange(height, device=grid_thw.device), + torch.arange(width, device=grid_thw.device), + indexing="ij", + ) + shape = (height // merge_size, merge_size, width // merge_size, merge_size) + hpos = hpos.reshape(shape).transpose(1, 2).flatten() + wpos = wpos.reshape(shape).transpose(1, 2).flatten() + positions.append(torch.stack((hpos, wpos), dim=-1).repeat(temporal, 1)) + return torch.cat(positions, dim=0) + + +def _vision_cu_seqlens(grid_thw: torch.Tensor) -> torch.Tensor: + """Return per-frame attention boundaries ``[segments + 1]``.""" + lengths = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]) + return F.pad(lengths.cumsum(0, dtype=torch.int32), (1, 0)) + + +class Glm5NextVisionPatchEmbed(nn.Module): + """Conv3d patch projection from ``[patches, C*T*P*P]`` to vision hidden.""" + + def __init__(self, config: Glm5NextVisionConfig, dtype: torch.dtype) -> None: + super().__init__() + self.in_channels = config.in_channels + self.temporal_patch_size = config.temporal_patch_size + self.patch_size = config.patch_size + kernel = (self.temporal_patch_size, self.patch_size, self.patch_size) + self.proj = nn.Conv3d( + self.in_channels, + config.hidden_size, + kernel_size=kernel, + stride=kernel, + bias=True, + dtype=dtype, + ) + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + """Project flattened patches to ``[vision_tokens, vision_hidden]``.""" + patches = pixel_values.view( + -1, + self.in_channels, + self.temporal_patch_size, + self.patch_size, + self.patch_size, + ) + return self.proj(patches.to(self.proj.weight.dtype)).view(patches.shape[0], -1) + + +class Glm5NextVisionRotaryEmbedding(nn.Module): + """Two-axis rotary frequencies used by the vision transformer.""" + + def __init__(self, dim: int, theta: float = 10000.0) -> None: + super().__init__() + inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward(self, position_ids: torch.Tensor) -> torch.Tensor: + """Map H/W ids ``[tokens, 2]`` to frequencies ``[tokens, head_dim/2]``.""" + return (position_ids.unsqueeze(-1) * self.inv_freq).flatten(1) + + +class Glm5NextVisionAttention(nn.Module): + """Bidirectional per-image attention over ``[vision_tokens, hidden]``.""" + + def __init__(self, config: Glm5NextVisionConfig, dtype: torch.dtype) -> None: + super().__init__() + self.num_heads = config.num_heads + self.head_dim = config.hidden_size // config.num_heads + self.scaling = self.head_dim**-0.5 + self.dropout = config.attention_dropout + self.qkv = nn.Linear( + config.hidden_size, + 3 * config.hidden_size, + bias=config.attention_bias, + dtype=dtype, + ) + self.proj = nn.Linear( + config.hidden_size, + config.hidden_size, + bias=config.attention_bias, + dtype=dtype, + ) + self.q_norm = Glm5NextRMSNorm(self.head_dim, config.rms_norm_eps, dtype) + self.k_norm = Glm5NextRMSNorm(self.head_dim, config.rms_norm_eps, dtype) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + """Attend within each cu-seqlens segment and return ``[tokens, hidden]``.""" + sequence = hidden_states.shape[0] + query, key, value = self.qkv(hidden_states).view(sequence, 3, self.num_heads, self.head_dim).unbind(1) + query, key = self.q_norm(query), self.k_norm(key) + cos, sin = (item.unsqueeze(-2).float() for item in position_embeddings) + query_float, key_float = query.float(), key.float() + query = (query_float * cos + _rotate_half(query_float) * sin).to(query.dtype) + key = (key_float * cos + _rotate_half(key_float) * sin).to(key.dtype) + query, key, value = (item.transpose(0, 1).unsqueeze(0) for item in (query, key, value)) + lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist() + q_chunks, k_chunks, v_chunks = (torch.split(item, lengths, dim=2) for item in (query, key, value)) + output = [ + F.scaled_dot_product_attention( + q, + k, + v, + dropout_p=self.dropout if self.training else 0.0, + scale=self.scaling, + ) + for q, k, v in zip(q_chunks, k_chunks, v_chunks) + ] + output = torch.cat(output, dim=2).squeeze(0).transpose(0, 1).reshape(sequence, -1) + return self.proj(output) + + +class Glm5NextVisionMLP(nn.Module): + """Clamped SwiGLU vision feed-forward block.""" + + def __init__(self, config: Glm5NextVisionConfig, dtype: torch.dtype) -> None: + super().__init__() + self.swiglu_limit = config.swiglu_limit + self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=True, dtype=dtype) + self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=True, dtype=dtype) + self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=True, dtype=dtype) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Transform ``[tokens, vision_hidden]`` with clamped SwiGLU.""" + gate = self.gate_proj(hidden_states).clamp(max=self.swiglu_limit) + up = self.up_proj(hidden_states).clamp(-self.swiglu_limit, self.swiglu_limit) + return self.down_proj(F.silu(gate) * up) + + +class Glm5NextVisionBlock(nn.Module): + """Pre-norm bidirectional vision transformer block.""" + + def __init__(self, config: Glm5NextVisionConfig, dtype: torch.dtype) -> None: + super().__init__() + self.norm1 = Glm5NextRMSNorm(config.hidden_size, config.rms_norm_eps, dtype) + self.norm2 = Glm5NextRMSNorm(config.hidden_size, config.rms_norm_eps, dtype) + self.attn = Glm5NextVisionAttention(config, dtype) + self.mlp = Glm5NextVisionMLP(config, dtype) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + """Transform ``[vision_tokens, hidden]`` without cross-image attention.""" + hidden_states = hidden_states + self.attn(self.norm1(hidden_states), cu_seqlens, position_embeddings) + return hidden_states + self.mlp(self.norm2(hidden_states)) + + +class Glm5NextVisionPatchMerger(nn.Module): + """Post-downsample projection in the text hidden dimension.""" + + def __init__(self, config: Glm5NextVisionConfig, dtype: torch.dtype) -> None: + super().__init__() + dim = config.out_hidden_size + context = config.projection_intermediate_size + self.swiglu_limit = config.swiglu_limit + self.proj = nn.Linear(dim, dim, bias=False, dtype=dtype) + self.post_projection_norm = nn.LayerNorm(dim, dtype=dtype) + self.gate_proj = nn.Linear(dim, context, bias=False, dtype=dtype) + self.up_proj = nn.Linear(dim, context, bias=False, dtype=dtype) + self.down_proj = nn.Linear(context, dim, bias=False, dtype=dtype) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Project downsampled image tokens ``[tokens, text_hidden]``.""" + hidden_states = F.gelu(self.post_projection_norm(self.proj(hidden_states))) + gate = self.gate_proj(hidden_states).clamp(max=self.swiglu_limit) + up = self.up_proj(hidden_states).clamp(-self.swiglu_limit, self.swiglu_limit) + return self.down_proj(F.silu(gate) * up) + + +@dataclass +class Glm5NextVisionOutput: + """Raw and merged vision token representations.""" + + last_hidden_state: torch.Tensor + pooler_output: torch.Tensor + + +class Glm5NextVisionModel(nn.Module): + """Patch encoder returning image features in the language hidden size.""" + + def __init__(self, config: Glm5NextVisionConfig) -> None: + super().__init__() + self.config = config + self.spatial_merge_size = config.spatial_merge_size + dtype = get_dtype(getattr(config, "torch_dtype", getattr(config, "dtype", None)), torch.bfloat16) + self.patch_embed = Glm5NextVisionPatchEmbed(config, dtype) + head_dim = config.hidden_size // config.num_heads + self.rotary_pos_emb = Glm5NextVisionRotaryEmbedding(head_dim // 2) + self.blocks = nn.ModuleList([Glm5NextVisionBlock(config, dtype) for _ in range(config.depth)]) + self.post_layernorm = Glm5NextRMSNorm(config.hidden_size, config.rms_norm_eps, dtype) + self.downsample = nn.Conv2d( + config.hidden_size, + config.out_hidden_size, + kernel_size=config.spatial_merge_size, + stride=config.spatial_merge_size, + dtype=dtype, + ) + self.merger = Glm5NextVisionPatchMerger(config, dtype) + + @property + def dtype(self) -> torch.dtype: + """Return the patch embedding dtype.""" + return self.patch_embed.proj.weight.dtype + + def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> Glm5NextVisionOutput: + """Encode flattened patches using grid metadata ``[images, (t,h,w)]``.""" + positions = _vision_position_ids(grid_thw, self.spatial_merge_size) + cu_seqlens = _vision_cu_seqlens(grid_thw) + hidden_states = self.patch_embed(pixel_values) + rotary = self.rotary_pos_emb(positions) + embedding = torch.cat((rotary, rotary), dim=-1) + position_embeddings = (embedding.cos(), embedding.sin()) + for block in self.blocks: + hidden_states = block(hidden_states, cu_seqlens, position_embeddings) + hidden_states = self.post_layernorm(hidden_states) + merge = self.spatial_merge_size + downsampled = hidden_states.view(-1, merge, merge, hidden_states.shape[-1]).permute(0, 3, 1, 2) + downsampled = self.downsample(downsampled).view(-1, self.config.out_hidden_size) + return Glm5NextVisionOutput(last_hidden_state=downsampled, pooler_output=self.merger(downsampled)) + + @torch.no_grad() + def init_weights(self, buffer_device: torch.device, init_std: float) -> None: + """Initialize vision parameters without materializing outside ``buffer_device``.""" + with buffer_device: + for module in self.modules(): + if isinstance(module, (nn.Linear, nn.Conv2d, nn.Conv3d)): + nn.init.normal_(module.weight, mean=0.0, std=init_std) + if module.bias is not None: + module.bias.zero_() + elif isinstance(module, (nn.LayerNorm, Glm5NextRMSNorm)): + module.reset_parameters() + + +__all__ = ["Glm5NextVisionModel", "Glm5NextVisionOutput"] diff --git a/nemo_automodel/components/models/glm_moe_dsa/kernels/cudnn_dsa.py b/nemo_automodel/components/models/glm_moe_dsa/kernels/cudnn_dsa.py index 5a0e504ee5..1e025e9fdf 100644 --- a/nemo_automodel/components/models/glm_moe_dsa/kernels/cudnn_dsa.py +++ b/nemo_automodel/components/models/glm_moe_dsa/kernels/cudnn_dsa.py @@ -16,11 +16,13 @@ from __future__ import annotations -import math from dataclasses import dataclass import torch +from nemo_automodel.components.models.common.cudnn_sparse_attention import ( + cudnn_sparse_attention as _shared_cudnn_sparse_attention, +) from nemo_automodel.shared.import_utils import safe_import_from _HAS_CUDNN_DSA, _CUDNN_DSA = safe_import_from( @@ -39,9 +41,7 @@ _INDEX_HEAD_DIM = 128 _ATTENTION_HEAD_DIM = 576 -_VALUE_HEAD_DIM = 512 _MAX_TOPK = 2048 -_FLASH_MLA_TOPK_ALIGNMENT = 512 _TOPK_SCRATCH_LIMIT_BYTES = 2 * 1024 * 1024 * 1024 _TOPK_SCRATCH_INT32_FACTOR = 2 _TOPK_ROW_ALIGNMENT = 512 @@ -565,185 +565,6 @@ def cudnn_indexer_topk( return global_indices.unsqueeze(1).contiguous() -def _padded_head_count(num_heads: int, major: int) -> int: - """Return the FlashMLA-supported head count for one SM generation.""" - if major >= 10: - for padded in (64, 128): - if num_heads == padded or (num_heads < padded and padded % num_heads == 0): - return padded - alignment = 128 - else: - alignment = 64 - if num_heads % alignment == 0: - return num_heads - if num_heads < alignment and alignment % num_heads == 0: - return alignment - raise ValueError(f"FlashMLA sparse prefill requires the query-head count to divide {alignment}, got H={num_heads}.") - - -def _pad_attention_heads( - q: torch.Tensor, attn_sink: torch.Tensor, padded_heads: int -) -> tuple[torch.Tensor, torch.Tensor]: - """Pad query ``[T, H, 576]`` and sink ``[H]`` to ``padded_heads``.""" - if q.shape[1] == padded_heads: - return q, attn_sink - q_padded = q.new_zeros((q.shape[0], padded_heads, q.shape[2])) - q_padded[:, : q.shape[1]] = q - sink_padded = attn_sink.new_full((padded_heads,), float("-inf")) - sink_padded[: q.shape[1]] = attn_sink - return q_padded, sink_padded - - -class _CudnnSparseAttention(torch.autograd.Function): - """Pair FlashMLA forward with cuDNN backward for latent THD attention.""" - - @staticmethod - def forward( - ctx, - q: torch.Tensor, - kv_latent: torch.Tensor, - topk_indices: torch.Tensor, - softmax_scale: float, - padded_heads: int, - topk_length: torch.Tensor | None, - all_rows_nonempty: bool, - valid_row_indices: torch.Tensor | None, - ) -> torch.Tensor: - """Run FlashMLA forward and save tensors required by cuDNN backward. - - Args: - ctx: Autograd context used to save forward tensors and scalar metadata. - q: CUDA BF16 query tensor of shape ``[T_q, H, 576]``. - kv_latent: CUDA BF16 gathered K/V tensor of shape ``[T_k, 1, 576]``. - topk_indices: CUDA int32 tensor of shape ``[T_q, 1, K]`` containing - global padded-storage K/V coordinates and a ``-1`` suffix. - softmax_scale: Scale applied to query-key scores. - padded_heads: FlashMLA-compatible padded head count. - topk_length: Optional int32 valid-prefix lengths of shape ``[T_q]``. - all_rows_nonempty: Whether every query has a positive valid prefix. - valid_row_indices: Optional int64 indices of nonempty queries with shape - ``[T_valid]``. - - Returns: - CUDA BF16 latent values of shape ``[T_q, H, 512]``. - """ - kv = kv_latent.squeeze(1).contiguous() - if topk_length is None: - indices, topk_length = _compact_and_sort_indices(topk_indices.squeeze(1), kv.shape[0]) - else: - # Model-prepared lengths imply the indexer already emitted a compact, sorted - # valid prefix. Reuse it directly instead of sorting [T, K] in every shared layer. - indices = topk_indices.squeeze(1).contiguous() - padded_topk = math.ceil(indices.shape[-1] / _FLASH_MLA_TOPK_ALIGNMENT) * _FLASH_MLA_TOPK_ALIGNMENT - if padded_topk != indices.shape[-1]: - indices = torch.nn.functional.pad(indices, (0, padded_topk - indices.shape[-1]), value=-1) - - attn_sink = torch.full((q.shape[1],), float("-inf"), dtype=torch.float32, device=q.device) - q_kernel, sink_kernel = _pad_attention_heads(q.contiguous(), attn_sink, padded_heads) - out_kernel, _max_logits, lse_kernel = _FLASH_MLA_SPARSE_FWD( - q_kernel, - kv.unsqueeze(1), - indices.unsqueeze(1), - softmax_scale, - d_v=_VALUE_HEAD_DIM, - attn_sink=sink_kernel, - topk_length=topk_length, - indexer_topk=0, - ) - out = out_kernel[:, : q.shape[1]].contiguous() - lse = lse_kernel[:, : q.shape[1]].contiguous() - if not all_rows_nonempty: - out.masked_fill_(topk_length.eq(0).view(-1, 1, 1), 0) - cached_valid_rows = ( - valid_row_indices if valid_row_indices is not None else torch.empty(0, dtype=torch.int64, device=q.device) - ) - ctx.save_for_backward(q, kv, out, lse, attn_sink, indices.clamp_min(0), topk_length, cached_valid_rows) - ctx.softmax_scale = softmax_scale - ctx.padded_heads = padded_heads - ctx.all_rows_nonempty = all_rows_nonempty - ctx.has_cached_valid_rows = valid_row_indices is not None - return out - - @staticmethod - def backward(ctx, grad_output: torch.Tensor): - """Map output gradients to query and gathered latent-KV layouts. - - Args: - ctx: Autograd context populated by :meth:`forward`. - grad_output: CUDA BF16 output gradient of shape ``[T_q, H, 512]``. - - Returns: - Gradients for the eight forward inputs: query ``[T_q, H, 576]``, - gathered latent K/V ``[T_k, 1, 576]``, then ``None`` for the index - and metadata inputs. - """ - q, kv, out, lse, attn_sink, indices, topk_length, cached_valid_rows = ctx.saved_tensors - valid_row_indices = None - if not ctx.all_rows_nonempty: - valid_row_indices = ( - cached_valid_rows - if ctx.has_cached_valid_rows - else torch.nonzero(topk_length > 0, as_tuple=False).flatten() - ) - - q_input = q - out_input = out - grad_input = grad_output - lse_input = lse - indices_kernel = indices - topk_length_kernel = topk_length - used_dummy = False - if valid_row_indices is not None: - if valid_row_indices.numel() == 0: - used_dummy = True - q_input = torch.zeros_like(q[:1]) - out_input = torch.zeros_like(out[:1]) - grad_input = torch.zeros_like(grad_output[:1]) - lse_input = torch.zeros_like(lse[:1]) - indices_kernel = torch.zeros_like(indices[:1]) - topk_length_kernel = torch.ones_like(topk_length[:1]) - else: - q_input = q.index_select(0, valid_row_indices) - out_input = out.index_select(0, valid_row_indices) - grad_input = grad_output.index_select(0, valid_row_indices) - lse_input = lse.index_select(0, valid_row_indices) - indices_kernel = indices.index_select(0, valid_row_indices) - topk_length_kernel = topk_length.index_select(0, valid_row_indices) - - q_kernel, sink_kernel = _pad_attention_heads(q_input, attn_sink, ctx.padded_heads) - if ctx.padded_heads == q.shape[1]: - out_kernel = out_input - grad_kernel = grad_input.contiguous() - lse_kernel = lse_input - else: - out_kernel = out_input.new_zeros((out_input.shape[0], ctx.padded_heads, out_input.shape[2])) - out_kernel[:, : out_input.shape[1]] = out_input - grad_kernel = grad_input.new_zeros((grad_input.shape[0], ctx.padded_heads, grad_input.shape[2])) - grad_kernel[:, : grad_input.shape[1]] = grad_input - lse_kernel = lse_input.new_zeros((lse_input.shape[0], ctx.padded_heads)) - lse_kernel[:, : lse_input.shape[1]] = lse_input - - result = _CUDNN_DSA.sparse_attention_backward_wrapper( - q_kernel.contiguous(), - kv, - out_kernel.contiguous(), - grad_kernel.contiguous(), - lse_kernel.contiguous(), - sink_kernel, - indices_kernel, - softmax_scale=ctx.softmax_scale, - topk_length=topk_length_kernel, - ) - if valid_row_indices is None: - grad_q = result["dq"][:, : q.shape[1]].contiguous() - else: - grad_q_valid = result["dq"][:0, : q.shape[1]] if used_dummy else result["dq"][:, : q.shape[1]] - grad_q = torch.zeros_like(q) - grad_q.index_copy_(0, valid_row_indices, grad_q_valid) - grad_kv = result["dkv"].unsqueeze(1).contiguous() - return grad_q, grad_kv, None, None, None, None, None, None - - def cudnn_sparse_attention( q: torch.Tensor, kv_latent: torch.Tensor, @@ -753,90 +574,45 @@ def cudnn_sparse_attention( all_rows_nonempty: bool = False, valid_row_indices: torch.Tensor | None = None, ) -> torch.Tensor: - """Run split GLM-5.2 sparse MLA with FlashMLA forward and cuDNN backward. + """Run GLM-5.2 sparse MLA through the shared cuDNN adapter. Args: - q: Absorbed MLA query, CUDA BF16 THD ``[T_q, H, 576]``. The final - dimension is ``kv_lora_rank + qk_rope_head_dim`` (``512 + 64``). - kv_latent: Shared latent key/value, CUDA BF16 THD ``[T_kv, 1, 576]``. - topk_indices: Global padded-storage K/V indices, CUDA int32 - ``[T_q, 1, K]`` with ``-1`` for invalid slots. - softmax_scale: Already-computed MLA attention scale. It is forwarded - unchanged to both FlashMLA and cuDNN backward. - topk_length: Optional int32 valid-prefix lengths ``[T_q]`` prepared once - from packed causal metadata. When supplied, ``topk_indices`` must already - contain the indexer's canonical compact, ascending prefix. + q: Contiguous CUDA BF16 absorbed query tensor of shape + ``[query_tokens, heads, 576]``. + kv_latent: Contiguous CUDA BF16 latent K/V tensor of shape + ``[key_tokens, 1, 576]``. + topk_indices: Contiguous CUDA int32 tensor of shape + ``[query_tokens, 1, sparse_width]`` with global K/V coordinates + and invalid entries marked ``-1``. + softmax_scale: Scale forwarded unchanged to FlashMLA and cuDNN backward. + topk_length: Optional contiguous CUDA int32 valid-prefix lengths of shape + ``[query_tokens]``. all_rows_nonempty: Whether every query has a positive valid-prefix length. - The model supplies this cached metadata flag to keep unpadded inputs on - the allocation-free backward path. - valid_row_indices: Optional cached int64 row indices whose valid-prefix length - is positive. The model supplies these once per stage for padded inputs so - every attention layer can compact without rescanning CUDA metadata. + valid_row_indices: Optional contiguous CUDA int64 indices of nonempty + queries with shape ``[valid_query_tokens]``. Returns: - Latent sparse-attention output, CUDA BF16 ``[T_q, H, 512]``. The caller - applies the model-owned value up-projection ``w_vc``. - - Raises: - RuntimeError: If optional kernels, CUDA, or SM90+ are unavailable. - TypeError: If compute tensors are not BF16 or indices are not int32. - ValueError: If tensor layouts, dimensions, top-k, or scale are invalid. + Contiguous CUDA BF16 latent output tensor of shape + ``[query_tokens, heads, 512]``. """ - _require_available() - major, _ = _require_cuda_tensors("cuDNN DSA sparse attention", q, kv_latent, topk_indices) - if q.dtype != torch.bfloat16 or kv_latent.dtype != torch.bfloat16: - raise TypeError(f"q and kv_latent must be bfloat16, got {q.dtype} and {kv_latent.dtype}.") - if topk_indices.dtype != torch.int32: - raise TypeError(f"topk_indices must be int32, got {topk_indices.dtype}.") - if q.ndim != 3 or q.shape[-1] != _ATTENTION_HEAD_DIM: - raise ValueError(f"q must have shape [T_q, H, {_ATTENTION_HEAD_DIM}], got {tuple(q.shape)}.") - if kv_latent.ndim != 3 or kv_latent.shape[1:] != (1, _ATTENTION_HEAD_DIM): - raise ValueError(f"kv_latent must have shape [T_kv, 1, {_ATTENTION_HEAD_DIM}], got {tuple(kv_latent.shape)}.") - if topk_indices.ndim != 3 or topk_indices.shape[:2] != (q.shape[0], 1): - raise ValueError(f"topk_indices must have shape [T_q, 1, K], got {tuple(topk_indices.shape)}.") - _validate_topk(topk_indices.shape[-1]) - if kv_latent.shape[0] >= torch.iinfo(torch.int32).max: - raise ValueError("The flattened KV token count must fit in an int32 global index.") - if not isinstance(softmax_scale, (float, int)) or not math.isfinite(float(softmax_scale)): - raise TypeError("softmax_scale must be a finite Python float.") - if float(softmax_scale) <= 0.0: - raise ValueError(f"softmax_scale must be positive, got {softmax_scale}.") - if not isinstance(all_rows_nonempty, bool): - raise TypeError(f"all_rows_nonempty must be a bool, got {type(all_rows_nonempty).__name__}.") - if all_rows_nonempty and valid_row_indices is not None: - raise ValueError("valid_row_indices must be None when all_rows_nonempty is true.") - if topk_length is not None: - if topk_length.shape != (q.shape[0],) or topk_length.dtype != torch.int32 or topk_length.device != q.device: - raise ValueError( - "topk_length must be an int32 tensor on the query device with shape " - f"{(q.shape[0],)}, got shape={tuple(topk_length.shape)}, " - f"dtype={topk_length.dtype}, device={topk_length.device}." - ) - if not topk_length.is_contiguous(): - raise ValueError("topk_length must be contiguous.") - if valid_row_indices is not None: - if topk_length is None: - raise ValueError("valid_row_indices requires precomputed topk_length metadata.") - if ( - valid_row_indices.ndim != 1 - or valid_row_indices.dtype != torch.int64 - or valid_row_indices.device != q.device - or not valid_row_indices.is_contiguous() - ): - raise ValueError("valid_row_indices must be a contiguous int64 tensor on the query device.") - if valid_row_indices.numel() > q.shape[0]: - raise ValueError("valid_row_indices cannot contain more entries than query rows.") - - padded_heads = _padded_head_count(q.shape[1], major) - return _CudnnSparseAttention.apply( - q.contiguous(), - kv_latent.contiguous(), - topk_indices.contiguous(), - float(softmax_scale), - padded_heads, - topk_length, - all_rows_nonempty, - valid_row_indices, + if q.ndim == 3 and q.shape[-1] != _ATTENTION_HEAD_DIM: + raise ValueError( + f"GLM-5.2 q must have shape [query_tokens, heads, {_ATTENTION_HEAD_DIM}], got {tuple(q.shape)}." + ) + if kv_latent.ndim == 3 and kv_latent.shape[1:] != (1, _ATTENTION_HEAD_DIM): + raise ValueError( + f"GLM-5.2 kv_latent must have shape [key_tokens, 1, {_ATTENTION_HEAD_DIM}], got {tuple(kv_latent.shape)}." + ) + if topk_indices.ndim == 3: + _validate_topk(topk_indices.shape[-1]) + return _shared_cudnn_sparse_attention( + q, + kv_latent, + topk_indices, + softmax_scale, + topk_length=topk_length, + all_rows_nonempty=all_rows_nonempty, + valid_row_indices=valid_row_indices, ) diff --git a/tests/functional_tests/context_parallel/run_glm5_next_packed_cp_parity.py b/tests/functional_tests/context_parallel/run_glm5_next_packed_cp_parity.py new file mode 100644 index 0000000000..48a189b6cf --- /dev/null +++ b/tests/functional_tests/context_parallel/run_glm5_next_packed_cp_parity.py @@ -0,0 +1,264 @@ +#!/usr/bin/env python +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Eight-rank packed CP parity for GLM-5.3-Flash KDA plus KPool-DSA. + +The reference and distributed models receive the same packed documents. CP=8 +uses the production contiguous sharder: input ids stay global until embedding +and image-splice time, while auxiliary token fields are sharded immediately. +The test compares reconstructed logits, input-embedding gradients, and every +used parameter gradient against CP=1. + +Usage: + torchrun --standalone --nproc_per_node=8 \ + tests/functional_tests/context_parallel/run_glm5_next_packed_cp_parity.py +""" + +from __future__ import annotations + +import os +import sys + +import torch +import torch.distributed as dist +import torch.nn.functional as F + + +def _config(): + """Return a small BF16 hybrid model that still exercises real FLA kernels.""" + from nemo_automodel.components.models.glm5_next.config import ( + Glm5NextConfig, + Glm5NextTextConfig, + Glm5NextVisionConfig, + ) + + text = Glm5NextTextConfig( + vocab_size=96, + hidden_size=64, + intermediate_size=128, + moe_intermediate_size=32, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=4, + n_shared_experts=1, + n_routed_experts=8, + num_experts_per_tok=2, + kv_lora_rank=16, + q_lora_rank=32, + qk_rope_head_dim=0, + qk_nope_head_dim=16, + v_head_dim=16, + index_topk=16, + index_head_dim=16, + index_n_heads=4, + index_kpool=4, + linear_head_dim=16, + linear_num_heads=4, + linear_conv_kernel_dim=4, + hc_mult=2, + hc_sinkhorn_iters=3, + mlp_layer_types=["dense"] * 4, + layer_types=["linear_attention"] * 3 + ["deepseek_sparse_attention"], + indexer_types=["full"] * 4, + pad_token_id=0, + torch_dtype="bfloat16", + ) + vision = Glm5NextVisionConfig( + depth=1, + hidden_size=16, + num_heads=2, + patch_size=2, + temporal_patch_size=2, + spatial_merge_size=2, + out_hidden_size=64, + intermediate_size=32, + projection_intermediate_size=64, + torch_dtype="bfloat16", + ) + return Glm5NextConfig(text_config=text, vision_config=vision, image_token_id=95, pad_token_id=0) + + +def _model(device: torch.device): + from nemo_automodel.components.models.common import BackendConfig + from nemo_automodel.components.models.glm5_next.model import Glm5NextForConditionalGeneration + + backend = BackendConfig( + attn="sdpa", + linear="torch", + rms_norm="torch_fp32", + experts="torch", + dispatcher="torch", + rope_fusion=False, + enable_hf_state_dict_adapter=False, + ) + model = Glm5NextForConditionalGeneration(_config(), backend=backend).to(device) + model.initialize_weights(device, dtype=torch.bfloat16) + return model.train() + + +def _sync_model(model: torch.nn.Module) -> None: + for parameter in model.parameters(): + dist.broadcast(parameter.data, src=0) + for buffer in model.buffers(): + dist.broadcast(buffer.data, src=0) + + +def main() -> None: + if not {"RANK", "WORLD_SIZE", "LOCAL_RANK"}.issubset(os.environ): + print("ERROR: launch this script with torchrun.", file=sys.stderr) + sys.exit(1) + + dist.init_process_group("nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + if world_size != 8: + if rank == 0: + print(f"ERROR: this parity contract requires CP=8, got {world_size} ranks.", file=sys.stderr) + dist.destroy_process_group() + sys.exit(1) + + from torch.distributed.device_mesh import init_device_mesh + + from nemo_automodel.components.models.glm5_next.cp import shard_batch_for_glm5_next_cp + from nemo_automodel.components.moe.parallelizer import apply_cp + + torch.manual_seed(1234) + reference = _model(device) + distributed = _model(device) + distributed.load_state_dict(reference.state_dict()) + _sync_model(reference) + distributed.load_state_dict(reference.state_dict()) + + sequence = 128 + torch.manual_seed(4321) + input_ids = torch.randint(1, 95, (1, sequence), device=device) + targets = torch.roll(input_ids, shifts=-1, dims=1) + doc_ids = torch.tensor( + [[1] * 19 + [2] * 37 + [3] * 72], + dtype=torch.int32, + device=device, + ) + dist.broadcast(input_ids, src=0) + dist.broadcast(targets, src=0) + + ref_logits = reference(input_ids=input_ids, _packed_seq_ids=doc_ids).logits + ref_loss = F.cross_entropy(ref_logits.float().flatten(0, 1), targets.flatten()) + ref_loss.backward() + + cp_mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("cp",))["cp"] + apply_cp(distributed, cp_mesh) + + _, local_batch, _ = shard_batch_for_glm5_next_cp( + cp_mesh, + None, + { + "input_ids": input_ids.clone(), + "labels": input_ids.clone(), + "_packed_seq_ids": doc_ids.clone(), + }, + shard_primary=False, + ) + context = local_batch["glm5_next_packed_context"] + local_logits = distributed( + input_ids=local_batch["input_ids"], + padding_mask=local_batch["padding_mask"], + glm5_next_packed_context=context, + ).logits + local_start = context.seq_start + local_end = local_start + context.local_seq_len + local_loss = F.cross_entropy( + local_logits.float().flatten(0, 1), + targets[:, local_start:local_end].flatten(), + ) + (local_loss / world_size).backward() + cp_loss = local_loss.detach().clone() / world_size + dist.all_reduce(cp_loss, op=dist.ReduceOp.SUM) + torch.testing.assert_close(cp_loss, ref_loss.detach(), rtol=2e-3, atol=2e-3) + + gathered_logits = [torch.empty_like(local_logits) for _ in range(world_size)] + dist.all_gather(gathered_logits, local_logits.detach()) + cp_logits = torch.cat(gathered_logits, dim=1) + logit_diff = (cp_logits.float() - ref_logits.detach().float()).abs() + if rank == 0: + print( + f"GLM-5.3 packed CP logits mean/max={logit_diff.mean().item():.3e}/{logit_diff.max().item():.3e}", + flush=True, + ) + # FLA transports recurrent state rank-to-rank in CP=8, changing BF16 + # accumulation order relative to the monolithic CP=1 chunked kernel. + torch.testing.assert_close(cp_logits, ref_logits.detach(), rtol=1.5e-1, atol=1.5e-1) + + max_grad_abs = 0.0 + gradient_diff_sq = 0.0 + gradient_ref_sq = 0.0 + gradient_dot = 0.0 + gradient_cp_sq = 0.0 + per_parameter_relative_l2 = [] + compared = 0 + cp_parameters = dict(distributed.named_parameters()) + for name, ref_parameter in reference.named_parameters(): + cp_parameter = cp_parameters[name] + if ref_parameter.grad is None: + if cp_parameter.grad is not None: + raise AssertionError(f"CP produced an unexpected gradient for {name}") + continue + if cp_parameter.grad is None: + raise AssertionError(f"CP did not produce a gradient for {name}") + cp_gradient = cp_parameter.grad.detach().float().clone() + dist.all_reduce(cp_gradient, op=dist.ReduceOp.SUM) + ref_gradient = ref_parameter.grad.detach().float() + if not torch.isfinite(cp_gradient).all(): + raise AssertionError(f"CP produced a non-finite gradient for {name}") + difference = cp_gradient - ref_gradient + ref_sq = ref_gradient.double().square().sum().item() + diff_sq = difference.double().square().sum().item() + cp_sq = cp_gradient.double().square().sum().item() + max_grad_abs = max(max_grad_abs, difference.abs().max().item()) + gradient_diff_sq += diff_sq + gradient_ref_sq += ref_sq + gradient_cp_sq += cp_sq + gradient_dot += (cp_gradient.double() * ref_gradient.double()).sum().item() + relative_l2 = diff_sq**0.5 / max(ref_sq**0.5, 1e-12) + per_parameter_relative_l2.append((relative_l2, name)) + compared += 1 + + gradient_relative_l2 = (gradient_diff_sq / gradient_ref_sq) ** 0.5 + gradient_cosine = gradient_dot / max((gradient_cp_sq * gradient_ref_sq) ** 0.5, 1e-12) + if rank == 0: + worst_parameters = ", ".join( + f"{name}={relative_l2:.3e}" for relative_l2, name in sorted(per_parameter_relative_l2, reverse=True)[:5] + ) + print( + "GLM-5.3 packed CP parity PASS: " + f"CP1 vs CP8, loss={cp_loss.item():.6f}/{ref_loss.item():.6f}, " + f"logits mean/max={logit_diff.mean().item():.3e}/{logit_diff.max().item():.3e}, " + f"parameter gradients compared={compared}, grad rel-L2/cosine/max-abs=" + f"{gradient_relative_l2:.3e}/{gradient_cosine:.6f}/{max_grad_abs:.3e}; " + f"worst tensor rel-L2: {worst_parameters}" + ) + if gradient_relative_l2 >= 5e-2 or gradient_cosine <= 0.998: + raise AssertionError( + f"CP gradient parity failed: relative L2={gradient_relative_l2:.3e}, cosine={gradient_cosine:.6f}" + ) + + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/unit_tests/models/glm5_next/__init__.py b/tests/unit_tests/models/glm5_next/__init__.py new file mode 100644 index 0000000000..26496bfed7 --- /dev/null +++ b/tests/unit_tests/models/glm5_next/__init__.py @@ -0,0 +1 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. diff --git a/tests/unit_tests/models/glm5_next/conftest.py b/tests/unit_tests/models/glm5_next/conftest.py new file mode 100644 index 0000000000..0e5b840efd --- /dev/null +++ b/tests/unit_tests/models/glm5_next/conftest.py @@ -0,0 +1,76 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import torch + +from nemo_automodel.components.models.common import BackendConfig +from nemo_automodel.components.models.glm5_next.config import ( + Glm5NextConfig, + Glm5NextTextConfig, + Glm5NextVisionConfig, +) +from nemo_automodel.components.models.glm5_next.model import Glm5NextForConditionalGeneration + + +def tiny_glm5_next_config() -> Glm5NextConfig: + """Build a four-layer hybrid config with one sparse MoE/DSA layer.""" + text = Glm5NextTextConfig( + vocab_size=64, + hidden_size=16, + intermediate_size=32, + moe_intermediate_size=8, + num_hidden_layers=4, + num_attention_heads=2, + num_key_value_heads=2, + n_shared_experts=1, + n_routed_experts=4, + num_experts_per_tok=2, + kv_lora_rank=8, + q_lora_rank=8, + qk_rope_head_dim=0, + qk_nope_head_dim=4, + v_head_dim=4, + index_topk=4, + index_head_dim=4, + index_n_heads=2, + index_kpool=2, + linear_head_dim=4, + linear_num_heads=2, + linear_conv_kernel_dim=2, + hc_mult=2, + hc_sinkhorn_iters=3, + mlp_layer_types=["dense", "dense", "dense", "sparse"], + layer_types=["linear_attention", "linear_attention", "linear_attention", "deepseek_sparse_attention"], + pad_token_id=0, + torch_dtype="float32", + ) + vision = Glm5NextVisionConfig( + depth=1, + hidden_size=8, + num_heads=2, + patch_size=2, + temporal_patch_size=2, + spatial_merge_size=2, + out_hidden_size=16, + intermediate_size=16, + projection_intermediate_size=32, + torch_dtype="float32", + ) + return Glm5NextConfig(text_config=text, vision_config=vision, image_token_id=63, pad_token_id=0) + + +def tiny_backend(*, adapter: bool = True) -> BackendConfig: + return BackendConfig( + attn="sdpa", + linear="torch", + rms_norm="torch", + experts="torch", + dispatcher="torch", + rope_fusion=False, + enable_hf_state_dict_adapter=adapter, + ) + + +def tiny_glm5_next_model() -> Glm5NextForConditionalGeneration: + model = Glm5NextForConditionalGeneration(tiny_glm5_next_config(), backend=tiny_backend()) + model.initialize_weights(torch.device("cpu"), dtype=torch.float32) + return model diff --git a/tests/unit_tests/models/glm5_next/test_config_registry.py b/tests/unit_tests/models/glm5_next/test_config_registry.py new file mode 100644 index 0000000000..7b46d14c67 --- /dev/null +++ b/tests/unit_tests/models/glm5_next/test_config_registry.py @@ -0,0 +1,65 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import pytest + +from nemo_automodel._transformers.registry import ( + _CUSTOM_CONFIG_REGISTRATIONS, + MODEL_ARCH_MAPPING, + resolve_custom_config_cls, +) +from nemo_automodel.components.models.glm5_next.config import Glm5NextConfig, Glm5NextTextConfig +from nemo_automodel.components.models.glm5_next.model import Glm5NextForConditionalGeneration + + +def test_checkpoint_style_nested_config_resolves_hybrid_patterns(): + config = Glm5NextConfig.from_dict( + { + "model_type": "glm5_next", + "architectures": ["Glm5NextForConditionalGeneration"], + "text_config": { + "num_hidden_layers": 4, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "q_lora_rank": 8, + "qk_rope_head_dim": 0, + "qk_nope_head_dim": 4, + "v_head_dim": 4, + "index_topk": 8, + "index_kpool": 4, + "layer_types": ["linear_attention", "linear_attention", "linear_attention", "full_attention"], + "linear_attn_config": { + "head_dim": 4, + "num_heads": 2, + "short_conv_kernel_size": 2, + "gate_lower_bound": -5.0, + }, + }, + "vision_config": {"depth": 1}, + } + ) + + assert config.text_config.layer_types[-1] == "deepseek_sparse_attention" + assert config.text_config.linear_lower_bound == -5.0 + assert config.text_config.linear_conv_kernel_dim == 2 + + +def test_config_rejects_rope_and_invalid_kpool_contracts(): + with pytest.raises(ValueError, match="NoPE"): + Glm5NextTextConfig(qk_rope_head_dim=16) + with pytest.raises(ValueError, match="divide index_topk"): + Glm5NextTextConfig(index_topk=7, index_kpool=4) + + +def test_registry_resolves_native_config_and_model(): + assert MODEL_ARCH_MAPPING["Glm5NextForConditionalGeneration"] == ( + "nemo_automodel.components.models.glm5_next.model", + "Glm5NextForConditionalGeneration", + ) + assert _CUSTOM_CONFIG_REGISTRATIONS["glm5_next"] == ( + "nemo_automodel.components.models.glm5_next.config", + "Glm5NextConfig", + ) + assert resolve_custom_config_cls("glm5_next") is Glm5NextConfig + capabilities = Glm5NextForConditionalGeneration.ModelCapabilities() + assert capabilities.supports_ep and capabilities.supports_cp and capabilities.supports_thd + assert not capabilities.supports_tp and not capabilities.supports_pp diff --git a/tests/unit_tests/models/glm5_next/test_cp.py b/tests/unit_tests/models/glm5_next/test_cp.py new file mode 100644 index 0000000000..968aae4ad8 --- /dev/null +++ b/tests/unit_tests/models/glm5_next/test_cp.py @@ -0,0 +1,91 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import torch +from torch.distributed.utils import _apply_to_tensors + +from nemo_automodel.components.distributed.context_parallel.sharder import contiguous_local_indices +from nemo_automodel.components.models.glm5_next.cp import ( + Glm5NextPackedContext, + doc_ids_from_cu_seqlens, + doc_ids_from_seq_lens, + segment_cu_seqlens, + shard_batch_for_glm5_next_cp, +) +from tests.unit_tests.models.glm5_next.conftest import tiny_glm5_next_model + + +class _FakeCPMesh: + def __init__(self, size: int, rank: int) -> None: + self._size = size + self._rank = rank + + def size(self) -> int: + return self._size + + def get_local_rank(self) -> int: + return self._rank + + +def test_document_metadata_conversions_preserve_padding_boundaries(): + expected = [[1, 1, 1, 2, 2, 0, 0]] + assert doc_ids_from_seq_lens(torch.tensor([[3, 2, -1000]]), 7).tolist() == expected + assert doc_ids_from_cu_seqlens(torch.tensor([0, 3, 5, -1000]), 7).tolist() == expected + assert segment_cu_seqlens(torch.tensor(expected[0], dtype=torch.int32)).tolist() == [0, 3, 5, 7] + + +def test_packed_context_can_be_rebuilt_by_fsdp_input_transform(): + context = Glm5NextPackedContext(torch.tensor([[1, 1, 2, 2]], dtype=torch.int32)) + context.row_cu_seqlens(0) + + rebuilt = _apply_to_tensors(lambda tensor: tensor.clone(), context) + + assert isinstance(rebuilt, Glm5NextPackedContext) + assert rebuilt.doc_ids.tolist() == [[1, 1, 2, 2]] + assert rebuilt.row_cu_seqlens(0)[0].tolist() == [0, 2, 4] + + +def test_vlm_cp_sharder_keeps_ids_and_media_global_but_shards_labels(): + batch = { + "input_ids": torch.arange(8).unsqueeze(0), + "labels": torch.arange(8).unsqueeze(0), + "attention_mask": torch.tensor([[1, 1, 1, 1, 2, 2, 2, 2]], dtype=torch.int32), + "pixel_values": torch.randn(8, 24), + "image_grid_thw": torch.tensor([[1, 2, 4]]), + } + + _, local, layout = shard_batch_for_glm5_next_cp(_FakeCPMesh(2, 1), None, batch) + + assert local["input_ids"].tolist() == [list(range(8))] + assert local["labels"].tolist() == [[4, 5, 6, 7]] + assert local["pixel_values"].shape == (8, 24) + assert local["image_grid_thw"].tolist() == [[1, 2, 4]] + context = local["glm5_next_packed_context"] + assert isinstance(context, Glm5NextPackedContext) + assert context.seq_start == 4 + assert context.local_doc_ids.tolist() == [[2, 2, 2, 2]] + assert (layout.original_seq_len, layout.padded_seq_len) == (8, 8) + + +def test_cp_sharder_pads_labels_without_padding_global_primary_ids(): + batch = { + "input_ids": torch.arange(6).unsqueeze(0), + "labels": torch.arange(6).unsqueeze(0), + "attention_mask": torch.ones(1, 6, dtype=torch.int32), + } + + _, local, layout = shard_batch_for_glm5_next_cp(_FakeCPMesh(4, 3), None, batch) + + assert local["input_ids"].shape == (1, 6) + assert local["labels"].tolist() == [[-100, -100]] + assert local["padding_mask"].all() + assert local["glm5_next_packed_context"].doc_ids.shape == (1, 8) + assert (layout.original_seq_len, layout.padded_seq_len) == (6, 8) + + +def test_model_registers_contiguous_model_owned_cp_sharder(): + model = tiny_glm5_next_model() + sharder = model.prepare_model_inputs_for_cp({"input_ids": torch.arange(8).unsqueeze(0)})["cp_sharder"] + + assert sharder.local_token_global_indices is contiguous_local_indices + assert model._owns_cp_attention + assert model._owns_packed_attention diff --git a/tests/unit_tests/models/glm5_next/test_model.py b/tests/unit_tests/models/glm5_next/test_model.py new file mode 100644 index 0000000000..2c21595386 --- /dev/null +++ b/tests/unit_tests/models/glm5_next/test_model.py @@ -0,0 +1,258 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import pytest +import torch + +from nemo_automodel._transformers.capabilities import ModelSupports +from nemo_automodel.components.models.glm5_next import layers as glm5_next_layers +from nemo_automodel.components.models.glm5_next.cp import Glm5NextPackedContext +from nemo_automodel.components.models.glm5_next.layers import ( + Glm5NextLinearAttention, + Glm5NextSparseAttention, +) +from nemo_automodel.components.models.glm5_next.model import build_glm5_next_moe_config +from tests.unit_tests.models.glm5_next.conftest import tiny_backend, tiny_glm5_next_config, tiny_glm5_next_model + + +def _torch_sparse_latent_attention( + query: torch.Tensor, + latent_kv: torch.Tensor, + indices: torch.Tensor, + softmax_scale: float, + *, + all_rows_nonempty: bool, +) -> torch.Tensor: + """Reference absorbed sparse attention used by the GLM-5.3 dispatch test. + + Args: + query: Absorbed query tensor with shape ``[queries, heads, latent_dim]``. + latent_kv: Shared latent K/V tensor with shape ``[keys, 1, latent_dim]``. + indices: Document-local sparse indices with shape + ``[queries, 1, sparse_width]``. + softmax_scale: Scale applied to query-key scores. + all_rows_nonempty: Whether every query has at least one selected key. + + Returns: + Latent attention output with shape ``[queries, heads, latent_dim]``. + """ + assert all_rows_nonempty + rows = [] + for row in range(query.shape[0]): + selected = indices[row, 0] + selected = selected[(selected >= 0) & (selected < latent_kv.shape[0])].unique(sorted=True) + assert selected.numel() > 0 + keys = latent_kv.index_select(0, selected.long()).squeeze(1) + scores = torch.matmul(query[row].float(), keys.float().transpose(0, 1)) * softmax_scale + rows.append(torch.matmul(scores.softmax(dim=-1), keys.float()).to(query.dtype)) + return torch.stack(rows) + + +def test_tiny_hybrid_packed_forward_backward_is_finite(): + torch.manual_seed(7) + model = tiny_glm5_next_model().train() + input_ids = torch.tensor([[1, 2, 3, 4, 5, 6]]) + document_ids = torch.tensor([[1, 1, 1, 2, 2, 2]], dtype=torch.int32) + + logits = model(input_ids=input_ids, attention_mask=document_ids).logits + logits.square().mean().backward() + + assert logits.shape == (1, 6, 64) + assert torch.isfinite(logits).all() + assert model.model.language_model.layers["0"].self_attn.q_proj.weight.grad is not None + assert model.model.language_model.layers["3"].self_attn.q_a_proj.weight.grad is not None + + +def test_torch_kda_decay_scales_the_key_dimension(): + q = torch.tensor([[[[0.0, 0.0]], [[1.0, 0.0]]]]) + k = torch.tensor([[[[1.0, 0.0]], [[0.0, 0.0]]]]) + v = torch.tensor([[[[1.0, 2.0]], [[0.0, 0.0]]]]) + g = torch.tensor([[[[1.0, 1.0]], [[0.5, 0.25]]]]).log() + beta = torch.tensor([[[1.0], [0.0]]]) + + output = glm5_next_layers._torch_recurrent_kda(q, k, v, g, beta, cu_seqlens=None) + + expected = torch.tensor([0.5, 1.0]) / torch.sqrt(torch.tensor(2.0)) + torch.testing.assert_close(output[0, 1, 0], expected) + + +def test_text_config_exposes_hidden_states_for_fused_linear_ce(): + model = tiny_glm5_next_model().eval() + model.config.text_config.output_hidden_states = True + + with torch.inference_mode(): + output = model(input_ids=torch.tensor([[1, 2, 3]]), logits_to_keep=1) + + assert output.logits.shape == (1, 1, 64) + assert output.hidden_states.shape == (1, 3, 16) + + +def test_packed_documents_are_attention_isolated(): + torch.manual_seed(11) + model = tiny_glm5_next_model().eval() + input_ids = torch.tensor([[1, 2, 3, 4, 5, 6]]) + changed = torch.tensor([[1, 2, 3, 13, 14, 15]]) + document_ids = torch.tensor([[1, 1, 1, 2, 2, 2]], dtype=torch.int32) + + with torch.inference_mode(): + baseline = model(input_ids=input_ids, attention_mask=document_ids).logits + perturbed = model(input_ids=changed, attention_mask=document_ids).logits + + torch.testing.assert_close(baseline[:, :3], perturbed[:, :3], rtol=1e-5, atol=1e-6) + assert not torch.allclose(baseline[:, 3:], perturbed[:, 3:]) + + +def test_sparse_indexer_prepares_document_pools_once_across_query_chunks(): + model = tiny_glm5_next_model().eval() + sparse = model.model.language_model.layers["3"].self_attn + sparse.query_chunk_size = 2 + calls = 0 + + def count_key_projection(_module, _inputs, _output): + nonlocal calls + calls += 1 + + handle = sparse.indexer.wk.register_forward_hook(count_key_projection) + try: + with torch.inference_mode(): + output = sparse._forward_document(torch.randn(1, 6, 16), 0, 6) + finally: + handle.remove() + + assert output.shape == (1, 6, 16) + assert calls == 1 + + +def test_cudnn_sparse_attention_matches_sdpa_absorbed_math(monkeypatch): + """The cuDNN dispatch preserves GLM-5.3 forward and gradient math.""" + torch.manual_seed(19) + config = tiny_glm5_next_config().text_config + config.kv_lora_rank = 512 + sdpa_backend = tiny_backend() + cudnn_backend = tiny_backend() + cudnn_backend.attn = "cudnn" + sdpa = Glm5NextSparseAttention(config, layer_idx=3, backend=sdpa_backend) + cudnn = Glm5NextSparseAttention(config, layer_idx=3, backend=cudnn_backend) + sdpa.init_weights(torch.device("cpu"), init_std=0.02) + cudnn.load_state_dict(sdpa.state_dict()) + + monkeypatch.setattr(glm5_next_layers, "is_cudnn_sparse_attention_available", lambda: True) + monkeypatch.setattr(glm5_next_layers, "cudnn_sparse_attention", _torch_sparse_latent_attention) + + sdpa_input = torch.randn(1, 6, config.hidden_size, requires_grad=True) + cudnn_input = sdpa_input.detach().clone().requires_grad_(True) + sdpa_output = sdpa._forward_document(sdpa_input, 0, 6) + cudnn_output = cudnn._forward_document(cudnn_input, 0, 6) + upstream = torch.randn_like(sdpa_output) + (sdpa_output * upstream).sum().backward() + (cudnn_output * upstream).sum().backward() + + torch.testing.assert_close(cudnn_output, sdpa_output, rtol=2e-4, atol=2e-6) + torch.testing.assert_close(cudnn_input.grad, sdpa_input.grad, rtol=5e-4, atol=2e-6) + for name in ("q_b_proj.weight", "kv_a_proj_with_mqa.weight", "kv_b_proj.weight"): + sdpa_grad = dict(sdpa.named_parameters())[name].grad + cudnn_grad = dict(cudnn.named_parameters())[name].grad + torch.testing.assert_close(cudnn_grad, sdpa_grad, rtol=5e-4, atol=2e-6) + + +def test_cudnn_sparse_attention_rejects_dropout(): + config = tiny_glm5_next_config().text_config + config.kv_lora_rank = 512 + config.attention_dropout = 0.1 + backend = tiny_backend() + backend.attn = "cudnn" + + with pytest.raises(ValueError, match="does not support attention dropout"): + Glm5NextSparseAttention(config, layer_idx=3, backend=backend) + + +def test_sparse_attention_empty_cp_shard_keeps_gather_backward_live(monkeypatch): + sparse = tiny_glm5_next_model().model.language_model.layers["3"].self_attn + backward_calls = 0 + + class FakeGather(torch.autograd.Function): + @staticmethod + def forward(ctx, local_hidden): + ctx.local_length = local_hidden.shape[1] + return torch.cat((local_hidden, local_hidden), dim=1) + + @staticmethod + def backward(ctx, full_grad): + nonlocal backward_calls + backward_calls += 1 + return full_grad[:, : ctx.local_length] + full_grad[:, ctx.local_length :] + + class FakeCPMesh: + @staticmethod + def get_group(): + return None + + monkeypatch.setattr( + glm5_next_layers, + "all_gather_sequence", + lambda local_hidden, _group, dim=1: FakeGather.apply(local_hidden), + ) + sparse.setup_cp_attention(FakeCPMesh()) + local_hidden = torch.randn(1, 2, 16, requires_grad=True) + context = Glm5NextPackedContext( + doc_ids=torch.tensor([[1, 1, 0, 0]], dtype=torch.int32), + seq_start=2, + cp_size=2, + ) + + output = sparse(local_hidden, packed_context=context) + output.sum().backward() + + assert output.requires_grad + assert backward_calls == 1 + torch.testing.assert_close(local_hidden.grad, torch.zeros_like(local_hidden)) + + +def test_image_features_replace_exactly_the_placeholder_tokens(): + torch.manual_seed(17) + model = tiny_glm5_next_model().eval() + input_ids = torch.tensor([[1, 63, 2, 3]]) + pixel_values = torch.randn(4, 3 * 2 * 2 * 2) + grid_thw = torch.tensor([[1, 2, 2]]) + + with torch.inference_mode(): + features = model.get_image_features(pixel_values, grid_thw).pooler_output + embeddings = model._embed_and_splice(input_ids, pixel_values, grid_thw) + logits = model( + input_ids=input_ids, + attention_mask=torch.ones_like(input_ids), + pixel_values=pixel_values, + image_grid_thw=grid_thw, + ).logits + + assert features.shape == (1, 16) + torch.testing.assert_close(embeddings[:, 1], features) + assert logits.shape == (1, 4, 64) + assert torch.isfinite(logits).all() + + +def test_hybrid_layer_pattern_and_parallel_capabilities(): + model = tiny_glm5_next_model() + layers = model.model.language_model.layers + assert all(isinstance(layers[str(index)].self_attn, Glm5NextLinearAttention) for index in range(3)) + assert isinstance(layers["3"].self_attn, Glm5NextSparseAttention) + + supports = ModelSupports(model, None) + assert supports.supports_cp + assert supports.supports_sequence_packing + assert supports.supports_cp_with_sequence_packing + assert "cudnn" in model._packed_cp_attn_backends + + +def test_hyperconnection_fp32_parameters_have_a_dedicated_fsdp_holder(): + hyperconnection = tiny_glm5_next_model().model.language_model.layers["0"].attn_hc + + assert set(dict(hyperconnection.named_parameters(recurse=False))) == {"fn"} + assert set(dict(hyperconnection._fp32_params.named_parameters())) == {"base", "scale"} + assert all(parameter.dtype is torch.float32 for parameter in hyperconnection._fp32_params.parameters()) + + +def test_moe_router_correction_bias_only_controls_expert_selection(): + config = build_glm5_next_moe_config(tiny_glm5_next_config().text_config, torch.float32) + + assert config.score_func == "sigmoid_with_bias" + assert not config.router_weight_uses_score_correction_bias diff --git a/tests/unit_tests/models/glm5_next/test_processing.py b/tests/unit_tests/models/glm5_next/test_processing.py new file mode 100644 index 0000000000..4ad4aa772f --- /dev/null +++ b/tests/unit_tests/models/glm5_next/test_processing.py @@ -0,0 +1,57 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import subprocess +import sys + +import pytest + +from nemo_automodel.components.models.glm5_next.processing import ( + _MEDIA_REMINDER, + _enable_image_placeholders, +) + + +def test_training_template_enables_images_but_keeps_other_media_disabled(): + template = "prefix " + _MEDIA_REMINDER + " suffix" + + patched = _enable_image_placeholders(template) + + assert "<|begin_of_image|><|image|><|end_of_image|>" in patched + assert "media_type == 'image'" in patched + assert "unable to process this" in patched + + +def test_existing_multimodal_template_is_not_rewritten(): + template = "{{ '<|image|>' }}" + assert _enable_image_placeholders(template) == template + + +def test_unrecognized_text_only_template_fails_loudly(): + with pytest.raises(ValueError, match="no recognized media rendering branch"): + _enable_image_placeholders("{{ messages }}") + + +def test_processor_modules_import_without_torchvision(): + script = r""" +import importlib +from unittest import mock + +from nemo_automodel.shared.import_utils import is_unavailable + +real_import_module = importlib.import_module + +def import_without_torchvision(name, *args, **kwargs): + if name == "torchvision.transforms.v2.functional": + raise ModuleNotFoundError("No module named 'torchvision'", name="torchvision") + if name == "transformers.models.glm46v.video_processing_glm46v": + raise ModuleNotFoundError("No module named 'torchvision'", name="torchvision") + return real_import_module(name, *args, **kwargs) + +with mock.patch("importlib.import_module", side_effect=import_without_torchvision): + image_processing = real_import_module("nemo_automodel.components.models.glm5_next.image_processing") + processing = real_import_module("nemo_automodel.components.models.glm5_next.processing") + +assert is_unavailable(image_processing.tvF) +assert is_unavailable(processing.Glm46VVideoProcessor) +""" + subprocess.run([sys.executable, "-c", script], check=True) diff --git a/tests/unit_tests/models/glm5_next/test_state_dict_adapter.py b/tests/unit_tests/models/glm5_next/test_state_dict_adapter.py new file mode 100644 index 0000000000..98f2606c99 --- /dev/null +++ b/tests/unit_tests/models/glm5_next/test_state_dict_adapter.py @@ -0,0 +1,112 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import torch +from torch.distributed.tensor import Shard + +from nemo_automodel.components.models.glm5_next.state_dict_adapter import ( + _apply_local_block_scales, + _local_shard_offsets, + dequantize_block_fp8, +) +from tests.unit_tests.models.glm5_next.conftest import tiny_glm5_next_model + + +def test_native_hf_round_trip_preserves_every_tensor(): + torch.manual_seed(29) + model = tiny_glm5_next_model() + native = model.state_dict() + + hf_state = model.state_dict_adapter.to_hf(native) + restored = model.state_dict_adapter.from_hf(dict(hf_state)) + + assert restored.keys() == native.keys() + for key, value in native.items(): + torch.testing.assert_close(restored[key], value, rtol=0.0, atol=0.0) + + +def test_adapter_routes_flat_hyperconnection_and_kda_fp32_parameters(): + adapter = tiny_glm5_next_model().state_dict_adapter + hf_state = { + "model.language_model.layers.0.hc_attn_fn": torch.randn(8, 32), + "model.language_model.layers.0.hc_attn_base": torch.randn(8), + "model.language_model.layers.0.hc_attn_scale": torch.randn(3), + "model.language_model.layers.0.self_attn.A_log": torch.randn(2), + "model.language_model.layers.0.self_attn.dt_bias": torch.randn(8), + } + + native = adapter.from_hf(dict(hf_state)) + + assert native["model.language_model.layers.0.attn_hc.fn"].shape == (8, 32) + assert native["model.language_model.layers.0.attn_hc._fp32_params.base"].dtype is torch.float32 + assert native["model.language_model.layers.0.attn_hc._fp32_params.scale"].dtype is torch.float32 + assert native["model.language_model.layers.0.self_attn._fp32_params.A_log"].shape == (2,) + assert native["model.language_model.layers.0.self_attn._fp32_params.A_log"].dtype is torch.float32 + assert native["model.language_model.layers.0.self_attn._fp32_params.dt_bias"].dtype is torch.float32 + assert adapter.to_hf(native).keys() == hf_state.keys() + + +def test_quantized_load_plan_matches_sparse_but_not_linear_output_projection(): + model = tiny_glm5_next_model() + planned = model.state_dict_adapter.to_hf(model.state_dict(), quantization=True, for_checkpoint_load=True) + + linear_o = "model.language_model.layers.0.self_attn.o_proj.weight" + sparse_o = "model.language_model.layers.3.self_attn.o_proj.weight" + assert linear_o in planned and linear_o + "_scale_inv" not in planned + assert sparse_o in planned and sparse_o + "_scale_inv" in planned + + +def test_mtp_layer_is_dropped_on_load(): + model = tiny_glm5_next_model() + layer_limit = model.config.text_config.num_hidden_layers + state = { + "lm_head.weight": torch.randn(64, 16), + f"model.language_model.layers.{layer_limit}.enorm.weight": torch.randn(16), + } + + converted = model.state_dict_adapter.from_hf(state) + + assert converted.keys() == {"lm_head.weight"} + + +def test_block_fp8_dequantization_uses_each_128_square_scale(): + weight = torch.ones((129, 129), dtype=torch.float8_e4m3fn) + scale = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + + output = dequantize_block_fp8(weight, scale, dtype=torch.float32) + + assert output[0, 0] == 1 + assert output[0, 128] == 2 + assert output[128, 0] == 3 + assert output[128, 128] == 4 + + +def test_block_fp8_dequantization_respects_misaligned_dtensor_shard_offset(): + weight = torch.ones((86, 129), dtype=torch.float8_e4m3fn) + scale = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + + output = _apply_local_block_scales(weight, scale, (58, 0), torch.float32) + + assert output[0, 0] == 1 + assert output[0, 128] == 2 + assert output[69, 0] == 1 + assert output[70, 0] == 3 + assert output[85, 128] == 4 + + +def test_dtensor_shard_offset_uses_torch_uneven_chunk_layout(): + class FakeMesh: + def size(self, mesh_dim): + assert mesh_dim == 0 + return 144 + + def get_local_rank(self, mesh_dim): + assert mesh_dim == 0 + return 23 + + class FakeDTensor: + ndim = 2 + shape = (12288, 4096) + placements = (Shard(0),) + device_mesh = FakeMesh() + + assert _local_shard_offsets(FakeDTensor()) == (1978, 0) diff --git a/tests/unit_tests/models/glm_moe_dsa/test_glm_moe_dsa_cudnn.py b/tests/unit_tests/models/glm_moe_dsa/test_glm_moe_dsa_cudnn.py index 5253b5d5b7..7d0a6c023a 100644 --- a/tests/unit_tests/models/glm_moe_dsa/test_glm_moe_dsa_cudnn.py +++ b/tests/unit_tests/models/glm_moe_dsa/test_glm_moe_dsa_cudnn.py @@ -21,6 +21,7 @@ import torch from transformers.models.glm_moe_dsa.configuration_glm_moe_dsa import GlmMoeDsaConfig +import nemo_automodel.components.models.common.cudnn_sparse_attention as shared_cudnn from nemo_automodel.components.models.common import BackendConfig from nemo_automodel.components.models.glm_moe_dsa import layers as layer_mod from nemo_automodel.components.models.glm_moe_dsa.kernels import cudnn_dsa @@ -641,11 +642,11 @@ def test_cudnn_sparse_attention_dispatches_padded_forward_and_exact_backward( ) -> None: fake_flash_mla = _FakeFlashMla() fake_dsa = _FakeSparseDsa(fake_flash_mla) - monkeypatch.setattr(cudnn_dsa, "_HAS_CUDNN_DSA", True) - monkeypatch.setattr(cudnn_dsa, "_HAS_FLASH_MLA", True) - monkeypatch.setattr(cudnn_dsa, "_CUDNN_DSA", fake_dsa) - monkeypatch.setattr(cudnn_dsa, "_FLASH_MLA_SPARSE_FWD", fake_flash_mla) - monkeypatch.setattr(cudnn_dsa, "_require_cuda_tensors", _accept_cpu_tensors) + monkeypatch.setattr(shared_cudnn, "_HAS_CUDNN_DSA", True) + monkeypatch.setattr(shared_cudnn, "_HAS_FLASH_MLA", True) + monkeypatch.setattr(shared_cudnn, "_CUDNN_DSA", fake_dsa) + monkeypatch.setattr(shared_cudnn, "_FLASH_MLA_SPARSE_FWD", fake_flash_mla) + monkeypatch.setattr(shared_cudnn, "_require_cuda_tensors", _accept_cpu_tensors) q = torch.ones(SPARSE_TOKENS, SPARSE_HEADS, QK_DIM, dtype=torch.bfloat16, requires_grad=True) kv_latent = torch.ones(SPARSE_TOKENS, 1, QK_DIM, dtype=torch.bfloat16, requires_grad=True) @@ -674,6 +675,88 @@ def test_cudnn_sparse_attention_dispatches_padded_forward_and_exact_backward( assert fake_dsa.called +def test_shared_cudnn_sparse_attention_accepts_glm53_dimensions(monkeypatch: pytest.MonkeyPatch) -> None: + """The shared adapter accepts GLM-5.3's 512 latent dim and 2051 raw indices.""" + query_tokens = 2 + key_tokens = 2304 + sparse_width = 2051 + + def fake_flash_mla( + q: torch.Tensor, + kv_latent: torch.Tensor, + indices: torch.Tensor, + softmax_scale: float, + *, + d_v: int, + attn_sink: torch.Tensor, + topk_length: torch.Tensor, + indexer_topk: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Validate GLM-5.3 kernel layouts and return deterministic tensors. + + Args: + q: Padded absorbed queries with shape ``[queries, heads, 512]``. + kv_latent: Shared latent K/V with shape ``[keys, 1, 512]``. + indices: Aligned sparse indices with shape ``[queries, 1, 2560]``. + softmax_scale: Scale applied to attention scores. + d_v: FlashMLA latent value width. + attn_sink: Per-head attention sinks with shape ``[heads]``. + topk_length: Valid index counts with shape ``[queries]``. + indexer_topk: Number of positions selected inside FlashMLA. + + Returns: + Output, maximum-logit, and log-sum-exp tensors in FlashMLA layouts. + """ + assert q.shape == (query_tokens, 64, 512) + assert kv_latent.shape == (key_tokens, 1, 512) + assert indices.shape == (query_tokens, 1, 2560) + assert softmax_scale == 0.125 + assert d_v == 512 + assert attn_sink.shape == (64,) + torch.testing.assert_close(topk_length, torch.full((query_tokens,), sparse_width, dtype=torch.int32)) + assert indexer_topk == 0 + output = torch.zeros(query_tokens, 64, 512, dtype=torch.bfloat16) + logits = torch.zeros(query_tokens, 64, dtype=torch.float32) + return output, logits, logits + + monkeypatch.setattr(shared_cudnn, "_HAS_CUDNN_DSA", True) + monkeypatch.setattr(shared_cudnn, "_HAS_FLASH_MLA", True) + monkeypatch.setattr(shared_cudnn, "_FLASH_MLA_SPARSE_FWD", fake_flash_mla) + monkeypatch.setattr(shared_cudnn, "_require_cuda_tensors", _accept_cpu_tensors) + + q = torch.ones(query_tokens, 64, 512, dtype=torch.bfloat16) + kv_latent = torch.ones(key_tokens, 1, 512, dtype=torch.bfloat16) + indices = torch.arange(sparse_width, dtype=torch.int32).view(1, 1, -1).expand(query_tokens, -1, -1) + output = shared_cudnn.cudnn_sparse_attention( + q, + kv_latent, + indices, + softmax_scale=0.125, + all_rows_nonempty=True, + ) + + assert output.shape == (query_tokens, 64, 512) + + +def test_glm52_sparse_wrapper_preserves_model_specific_dimensions() -> None: + """Sharing the adapter does not widen GLM-5.2's public tensor contract.""" + with pytest.raises(ValueError, match="GLM-5.2 q must have shape"): + cudnn_dsa.cudnn_sparse_attention( + torch.ones(2, 64, 512, dtype=torch.bfloat16), + torch.ones(2, 1, 512, dtype=torch.bfloat16), + torch.zeros(2, 1, 1, dtype=torch.int32), + softmax_scale=0.125, + ) + + with pytest.raises(ValueError, match="index_topk must be in"): + cudnn_dsa.cudnn_sparse_attention( + torch.ones(2, 64, 576, dtype=torch.bfloat16), + torch.ones(2, 1, 576, dtype=torch.bfloat16), + torch.zeros(2, 1, 2049, dtype=torch.int32), + softmax_scale=0.125, + ) + + @pytest.mark.parametrize("cache_valid_rows", [False, True]) def test_cudnn_sparse_attention_compacts_zero_length_rows( monkeypatch: pytest.MonkeyPatch, cache_valid_rows: bool @@ -682,11 +765,11 @@ def test_cudnn_sparse_attention_compacts_zero_length_rows( topk_length = torch.tensor([0, 2], dtype=torch.int32) fake_flash_mla = _FakeZeroLengthFlashMla(topk_length) fake_dsa = _FakeCompactedSparseDsa(valid_rows=1, dkv_value=5) - monkeypatch.setattr(cudnn_dsa, "_HAS_CUDNN_DSA", True) - monkeypatch.setattr(cudnn_dsa, "_HAS_FLASH_MLA", True) - monkeypatch.setattr(cudnn_dsa, "_CUDNN_DSA", fake_dsa) - monkeypatch.setattr(cudnn_dsa, "_FLASH_MLA_SPARSE_FWD", fake_flash_mla) - monkeypatch.setattr(cudnn_dsa, "_require_cuda_tensors", _accept_cpu_tensors) + monkeypatch.setattr(shared_cudnn, "_HAS_CUDNN_DSA", True) + monkeypatch.setattr(shared_cudnn, "_HAS_FLASH_MLA", True) + monkeypatch.setattr(shared_cudnn, "_CUDNN_DSA", fake_dsa) + monkeypatch.setattr(shared_cudnn, "_FLASH_MLA_SPARSE_FWD", fake_flash_mla) + monkeypatch.setattr(shared_cudnn, "_require_cuda_tensors", _accept_cpu_tensors) q = torch.ones(SPARSE_TOKENS, SPARSE_HEADS, QK_DIM, dtype=torch.bfloat16, requires_grad=True) kv_latent = torch.ones(SPARSE_TOKENS, 1, QK_DIM, dtype=torch.bfloat16, requires_grad=True) @@ -718,11 +801,11 @@ def test_cudnn_sparse_attention_all_empty_rows_use_only_dummy(monkeypatch: pytes topk_length = torch.zeros(SPARSE_TOKENS, dtype=torch.int32) fake_flash_mla = _FakeZeroLengthFlashMla(topk_length) fake_dsa = _FakeCompactedSparseDsa(valid_rows=0, dkv_value=0) - monkeypatch.setattr(cudnn_dsa, "_HAS_CUDNN_DSA", True) - monkeypatch.setattr(cudnn_dsa, "_HAS_FLASH_MLA", True) - monkeypatch.setattr(cudnn_dsa, "_CUDNN_DSA", fake_dsa) - monkeypatch.setattr(cudnn_dsa, "_FLASH_MLA_SPARSE_FWD", fake_flash_mla) - monkeypatch.setattr(cudnn_dsa, "_require_cuda_tensors", _accept_cpu_tensors) + monkeypatch.setattr(shared_cudnn, "_HAS_CUDNN_DSA", True) + monkeypatch.setattr(shared_cudnn, "_HAS_FLASH_MLA", True) + monkeypatch.setattr(shared_cudnn, "_CUDNN_DSA", fake_dsa) + monkeypatch.setattr(shared_cudnn, "_FLASH_MLA_SPARSE_FWD", fake_flash_mla) + monkeypatch.setattr(shared_cudnn, "_require_cuda_tensors", _accept_cpu_tensors) q = torch.ones(SPARSE_TOKENS, SPARSE_HEADS, QK_DIM, dtype=torch.bfloat16, requires_grad=True) kv_latent = torch.ones(SPARSE_TOKENS, 1, QK_DIM, dtype=torch.bfloat16, requires_grad=True) diff --git a/tests/unit_tests/recipes/test_glm5_next_medpix_recipes.py b/tests/unit_tests/recipes/test_glm5_next_medpix_recipes.py new file mode 100644 index 0000000000..c2a4329a50 --- /dev/null +++ b/tests/unit_tests/recipes/test_glm5_next_medpix_recipes.py @@ -0,0 +1,48 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from __future__ import annotations + +from pathlib import Path + +import yaml + +_RECIPES = Path("examples/vlm_finetune/glm5_next") +_RECIPE_NAME = "glm5_3_flash_medpix_packed2k_ep72_cp2_100steps.yaml" + + +def _load() -> dict: + return yaml.safe_load((_RECIPES / _RECIPE_NAME).read_text()) + + +def test_glm5_next_medpix_ep72_cp2_recipe_contract(): + recipe = _load() + + assert recipe["recipe"] == "FinetuneRecipeForVLM" + assert recipe["model"]["pretrained_model_name_or_path"] == "zai-org/GLM-5.3-Flash" + assert recipe["model"]["backend"]["attn"] == "cudnn" + assert recipe["processor"]["_target_"].endswith("build_glm5_next_processor") + assert recipe["dataset"] == { + "_target_": "nemo_automodel.components.datasets.vlm.datasets.make_medpix_dataset", + "path_or_dataset": "mmoukouba/MedPix-VQA", + "split": "train", + } + assert recipe["step_scheduler"]["max_steps"] == 100 + assert recipe["step_scheduler"]["global_batch_size"] == 144 + assert recipe["step_scheduler"]["local_batch_size"] == 1 + assert recipe["distributed"]["strategy"] == "fsdp2" + assert recipe["distributed"]["ep_size"] == 72 + assert recipe["distributed"]["cp_size"] == 2 + assert recipe["distributed"]["pp_size"] == 1 + assert recipe["distributed"]["tp_size"] == 1 + assert recipe["distributed"]["defer_fsdp_grad_sync"] is False + assert recipe["distributed"]["moe"]["wrap_outer_model"] is True + assert recipe["packed_sequence"]["packing_format"] == "neat" + assert recipe["packed_sequence"]["max_length"] == 2048 + assert recipe["packed_sequence"]["pack_size"] == 2048 + assert recipe["packed_sequence"]["collate_max_length"] == 2048 + assert recipe["wandb"]["enable"] is False + assert recipe["wandb"]["name"] == "glm5_3_flash_medpix_packed2k_ep72_cp2_100steps" + assert "ep72" in recipe["wandb"]["tags"] + assert "cp2" in recipe["wandb"]["tags"] + assert recipe["ci"]["nodes"] == 9 + assert recipe["ci"]["time"] == "01:00:00"