ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsGlm Moe Dsa

nemo_automodel.components.models.glm_moe_dsa.state_dict_adapter

View as Markdown

Module Contents

Classes

NameDescription
GlmMoeDsaStateDictAdapterConverts between HF GLM-MoE-DSA checkpoints and native format.

Functions

NameDescription
_dequantize_glm_fp8Dequantize GLM FP8 weights without changing other model adapters.
_dequantize_glm_with_torch_offsetsDequantize a GLM local shard while preserving the global FP8 block grid.
_dequantize_glm_with_triton_offsetsRun offset-aware GLM FP8 dequantization with Triton.
_glm_dtensor_local_offsetsReturn the exact global row/column origin of a GLM DTensor shard.
_glm_weight_dequant_offset_kernelDequantize a local GLM shard whose origin is inside a global FP8 block.
_slice_glm_scale_for_dtensorSlice the global GLM scale grid for one rank-local weight shard.
should_quantize_keyReturn whether key has a blockwise-FP8 scale in GLM checkpoints.

Data

NON_QUANTIZED_KEY_PATTERNS

_GLM_TRITON_AVAILABLE

logger

API

class nemo_automodel.components.models.glm_moe_dsa.state_dict_adapter.GlmMoeDsaStateDictAdapter()

Bases: Glm4MoeStateDictAdapter

Converts between HF GLM-MoE-DSA checkpoints and native format.

Extends Glm4MoeStateDictAdapter with handling for the DSA indexer weights that should not be quantized (k_norm, weights_proj).

_indexer_non_quantized_keys
nemo_automodel.components.models.glm_moe_dsa.state_dict_adapter.GlmMoeDsaStateDictAdapter._dequantize(
    state_dict: dict[str, typing.Any]
) -> dict[str, typing.Any]
nemo_automodel.components.models.glm_moe_dsa.state_dict_adapter.GlmMoeDsaStateDictAdapter._uses_blockwise_fp8_checkpoint() -> bool
nemo_automodel.components.models.glm_moe_dsa.state_dict_adapter.GlmMoeDsaStateDictAdapter.convert_single_tensor_to_hf(
    fqn: str,
    tensor: typing.Any,
    kwargs = {}
) -> list[tuple[str, typing.Any]]
nemo_automodel.components.models.glm_moe_dsa.state_dict_adapter.GlmMoeDsaStateDictAdapter.from_hf(
    hf_state_dict: dict[str, typing.Any],
    device_mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh] = None,
    kwargs = {}
) -> dict[str, typing.Any]

Dequantize blockwise FP8 tensors before merging per-expert weights.

nemo_automodel.components.models.glm_moe_dsa.state_dict_adapter._dequantize_glm_fp8(
    weight: torch.Tensor,
    scale_inv: torch.Tensor,
    dtype: torch.dtype = torch.bfloat16,
    block_size: int = BLOCK_SIZE,
    name: str = ''
) -> torch.Tensor

Dequantize GLM FP8 weights without changing other model adapters.

Parameters:

weight
torch.Tensor

FP8 matrix or DTensor matrix. A DTensor uses its global shape while to_local() supplies the rank-local [rows, cols] shard.

scale_inv
torch.Tensor

Matching blockwise scales. For a DTensor weight this is the full global scale grid loaded from the Hugging Face checkpoint.

dtype
torch.dtypeDefaults to torch.bfloat16

Output floating-point dtype.

block_size
intDefaults to BLOCK_SIZE

Height and width of each square FP8 quantization block.

name
strDefaults to ''

Checkpoint tensor name used in diagnostics.

Returns: torch.Tensor

Dequantized tensor preserving the input DTensor mesh and placements when present.

nemo_automodel.components.models.glm_moe_dsa.state_dict_adapter._dequantize_glm_with_torch_offsets(
    weight: torch.Tensor,
    scale_inv: torch.Tensor,
    dtype: torch.dtype,
    block_size: int,
    offsets_within_first_block: tuple[int, int]
) -> torch.Tensor

Dequantize a GLM local shard while preserving the global FP8 block grid.

Parameters:

weight
torch.Tensor

Local FP8 weight shard with shape [rows, cols].

scale_inv
torch.Tensor

Sliced scale grid with shape [block_rows, block_cols], where each dimension covers the local shard plus its offset into the first block.

dtype
torch.dtype

Output floating-point dtype.

block_size
int

Height and width of each square FP8 quantization block.

offsets_within_first_block
tuple[int, int]

Element offset (row, col) of the shard origin within its first global block. Each value is in [0, block_size).

Returns: torch.Tensor

Dequantized local tensor with the same shape as weight.

nemo_automodel.components.models.glm_moe_dsa.state_dict_adapter._dequantize_glm_with_triton_offsets(
    weight: torch.Tensor,
    scale_inv: torch.Tensor,
    dtype: torch.dtype,
    block_size: int,
    offsets_within_first_block: tuple[int, int]
) -> torch.Tensor

Run offset-aware GLM FP8 dequantization with Triton.

Parameters:

weight
torch.Tensor

Local CUDA FP8 weight shard with shape [rows, cols].

scale_inv
torch.Tensor

Sliced CUDA scale grid covering the shard’s intersecting global blocks.

dtype
torch.dtype

Output floating-point dtype.

block_size
int

Height and width of each square FP8 quantization block.

offsets_within_first_block
tuple[int, int]

Element offset (row, col) of the shard origin within its first global block. Each value is in [0, block_size).

Returns: torch.Tensor

Dequantized local CUDA tensor with the same shape as weight.

nemo_automodel.components.models.glm_moe_dsa.state_dict_adapter._glm_dtensor_local_offsets(
    weight_dtensor: torch.Tensor,
    weight_local: torch.Tensor,
    scale_inv: torch.Tensor,
    block_size: int
) -> tuple[int, int]

Return the exact global row/column origin of a GLM DTensor shard.

Parameters:

weight_dtensor
torch.Tensor

Global two-dimensional DTensor checkpoint destination.

weight_local
torch.Tensor

Rank-local weight shard with shape [rows, cols].

scale_inv
torch.Tensor

Full global blockwise scale grid. Used only by the fallback placement calculation when checkpoint chunk metadata is unavailable.

block_size
int

Height and width of each square FP8 quantization block.

Returns: tuple[int, int]

Absolute element offset (row, col) of the local shard in the global weight.

nemo_automodel.components.models.glm_moe_dsa.state_dict_adapter._glm_weight_dequant_offset_kernel(
    x_ptr,
    s_ptr,
    y_ptr,
    M,
    N,
    row_offset,
    col_offset,
    stride_xm,
    stride_xn,
    stride_ym,
    stride_yn,
    stride_sm,
    stride_sn,
    BLOCK_SIZE: triton.language.constexpr
)

Dequantize a local GLM shard whose origin is inside a global FP8 block.

nemo_automodel.components.models.glm_moe_dsa.state_dict_adapter._slice_glm_scale_for_dtensor(
    scale_inv: torch.Tensor,
    weight_dtensor: torch.Tensor,
    weight_local: torch.Tensor,
    block_size: int = BLOCK_SIZE
) -> torch.Tensor

Slice the global GLM scale grid for one rank-local weight shard.

Parameters:

scale_inv
torch.Tensor

Full global scale grid with shape [global_block_rows, global_block_cols].

weight_dtensor
torch.Tensor

Global two-dimensional DTensor checkpoint destination.

weight_local
torch.Tensor

Rank-local FP8 weight shard with shape [rows, cols].

block_size
intDefaults to BLOCK_SIZE

Height and width of each square FP8 quantization block.

Returns: torch.Tensor

Contiguous scale grid covering every global block intersected by the local shard.

nemo_automodel.components.models.glm_moe_dsa.state_dict_adapter.should_quantize_key(
    key: str
) -> bool

Return whether key has a blockwise-FP8 scale in GLM checkpoints.

nemo_automodel.components.models.glm_moe_dsa.state_dict_adapter.NON_QUANTIZED_KEY_PATTERNS = ['norm.weight', 'lm_head.weight', 'embed_tokens.weight', 'mlp.gate.weight', 'eh_...
 
nemo_automodel.components.models.glm_moe_dsa.state_dict_adapter._GLM_TRITON_AVAILABLE = True
 
nemo_automodel.components.models.glm_moe_dsa.state_dict_adapter.logger = logging.getLogger(__name__)