[ckpt] feat: support HF INT4 compressed-tensors dequantization

.github/workflows/secrets.yml

@@ -23,3 +23,4 @@ jobs:
       - uses: gitleaks/gitleaks-action@v2
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          GITLEAKS_LICENSE: ${{ secrets.GITLEAKS_LICENSE }}

configs/models/kimi_k2/ckpt_convert/kimi_k2_convert.yaml

@@ -28,10 +28,14 @@ name_map:
     transformer: language_model.model
     layer_prefix: layers
     input_layernorm: input_layernorm
-    attention.q_down: self_attn.q_a_proj
+    attention.q_down:
+      name: self_attn.q_a_proj
+      dtype: bf16
     attention.q_up: self_attn.q_b_proj
     attention.q_up_layernorm: self_attn.q_a_layernorm
-    attention.kv_down: self_attn.kv_a_proj_with_mqa
+    attention.kv_down:
+      name: self_attn.kv_a_proj_with_mqa
+      dtype: bf16
     attention.kv_up: self_attn.kv_b_proj
     attention.kv_up_layernorm: self_attn.kv_a_layernorm
     attention.dense: self_attn.o_proj
@@ -60,6 +64,7 @@ name_map:
       extra: true
       fp8: true
       fp8_ignore_tp: true
+      dtype: fp8
     attention.q_up:
       name: self_attention.linear_q_up_proj
       extra: true
@@ -72,6 +77,7 @@ name_map:
       extra: true
       fp8: true
       fp8_ignore_tp: true
+      dtype: fp8
     attention.kv_up:
       name: self_attention.linear_kv_up_proj
       extra: true
@@ -111,4 +117,3 @@ name_map:
     word_embeddings_for_head: language_model.output_layer
     transformer_tpl: model%d
 torch_dtype: bfloat16
-

docker/Dockerfile

@@ -204,4 +204,4 @@ RUN rm -rf /tmp/* ~/.bash_history && \
     rm -f /workspace/README.md /workspace/license.txt && \
     rm -rf /workspace/docker-examples /workspace/tutorials /workspace/lerobot
 
-WORKDIR /workspace
+WORKDIR /workspace

examples/kimi_k2.5/checkpoint_convert/convert_kimi_k2.5_hf_to_mcore.sh

@@ -5,20 +5,32 @@ MEGATRON_PATH=${MEGATRON_PATH:-"/workspace/Loong-Megatron"}
 CONVERT_CHECKPOINT_PATH="${LOONGFORGE_PATH}/tools/convert_checkpoint"
 
 LOAD=/mnt/cluster/huggingface.co/moonshotai/Kimi-K2.5/
-SAVE=/mnt/cluster/LoongForge/moonshotai/Kimi-K2.5-entp8dtp8pp8ep32etp1
+SAVE=/mnt/cluster/LoongForge/moonshotai/Kimi-K2.5-entp1dtp8pp8ep16etp1
 
 MODEL_CONFIG_FILE=${LOONGFORGE_PATH}/configs/models/kimi_k2.5/kimi_k2_5.yaml
 
 FOUNDATION_CONVERT_FILE=${LOONGFORGE_PATH}/configs/models/kimi_k2/ckpt_convert/kimi_k2_convert.yaml
 IMAGE_ENCODER_CONVERT_FILE=${LOONGFORGE_PATH}/configs/models/image_encoder/ckpt_convert/moon_vit_3d_convert.yaml
 IMAGE_PROJECTOR_CONVERT_FILE=${LOONGFORGE_PATH}/configs/models/image_projector/ckpt_convert/patch_merger_adapter_convert.yaml
 
-ETP=8
+ETP=1
 DTP=8
 PP=8
-EP=32
+EP=16
 Expert_TP=1
 
+EXTRA_ARGS=(
+    --hf-dequantize-int4
+    --hf-dequantize-dtype "${HF_DEQUANTIZE_DTYPE:-bfloat16}"
+)
+
+QUANT_CONFIG_FILE="${LOAD%/}/config.json"
+if [[ -f "$QUANT_CONFIG_FILE" ]]; then
+    EXTRA_ARGS+=(--hf-quant-config-file "$QUANT_CONFIG_FILE")
+else
+    echo "WARNING: compressed-tensors config not found: $QUANT_CONFIG_FILE; using fallback Kimi INT4 quantization args."
+fi
+
 
 PYTHONPATH=$MEGATRON_PATH:$PYTHONPATH \
     python $CONVERT_CHECKPOINT_PATH/module_convertor/model.py \
@@ -40,4 +52,8 @@ PYTHONPATH=$MEGATRON_PATH:$PYTHONPATH \
     --moe-grouped-gemm \
     --safetensors \
     --no_save_optim \
-    --no_load_optim
+    --no_load_optim \
+    --custom_pipeline_layers 8,8,8,8,8,8,8,5 \
+    --force_pow_2_scales \
+    --max_workers 64 \
+    "${EXTRA_ARGS[@]}"

pyproject.toml

@@ -37,6 +37,7 @@ gpu = [
     "tokenizers==0.22.2",
     "sentencepiece>=0.2.0",
     "safetensors==0.4.5",
+    "compressed-tensors==0.15.0.1",
     "huggingface_hub==1.5.0",
     "tiktoken==0.5.2",
     # Training / Modeling utils

tools/convert_checkpoint/arguments.py

@@ -87,6 +87,19 @@ def _add_checkpoint_args(parser):
                        help='Run pretrain as fp8, only used for hf to mcore when '
                        'saved checkpoint is bf16 and pretrain as fp8')
 
+    group.add_argument('--hf-dequantize-int4', '--hf_dequantize_int4', dest='hf_dequantize_int4',
+                       action='store_true',
+                       help=('When loading a HuggingFace checkpoint, materialize compressed-tensors '
+                             'pack-quantized INT4 weights into floating-point .weight tensors before '
+                             'the common/mcore conversion.'))
+    group.add_argument('--hf-dequantize-dtype', '--hf_dequantize_dtype', dest='hf_dequantize_dtype',
+                       type=str, default='bfloat16',
+                       choices=['bf16', 'bfloat16', 'fp16', 'float16', 'fp32', 'float32'],
+                       help='Floating dtype used for on-the-fly HF INT4 dequantization.')
+    group.add_argument('--hf-quant-config-file', '--hf_quant_config_file', dest='hf_quant_config_file',
+                       type=str, default=None,
+                       help='Optional HF config.json path for compressed-tensors quantization_config.')
+
     group.add_argument('--fp8_quant_transfer_type', type=str, default="float32", choices=["float32", "bfloat16"],
                        help='The transfer dtype when convert from hf fp8 to mcore fp8')
     group.add_argument('--distributed_convert', action='store_true',

tools/convert_checkpoint/huggingface/compressed_tensors_dequant.py

@@ -0,0 +1,267 @@
+#!/usr/bin/env python3
+# Copyright 2026 The LoongForge Authors.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Utilities for materializing compressed-tensors pack-quantized HF weights."""
+
+import json
+import logging
+from pathlib import Path
+
+import torch
+
+
+LOGGER = logging.getLogger(__name__)
+
+WEIGHT_PACKED_KEY = "weight_packed"
+WEIGHT_SCALE_KEY = "weight_scale"
+WEIGHT_SHAPE_KEY = "weight_shape"
+WEIGHT_ZERO_POINT_KEY = "weight_zero_point"
+WEIGHT_G_IDX_KEY = "weight_g_idx"
+
+DTYPE_MAP = {
+    "bf16": torch.bfloat16,
+    "bfloat16": torch.bfloat16,
+    "fp16": torch.float16,
+    "float16": torch.float16,
+    "fp32": torch.float32,
+    "float32": torch.float32,
+}
+
+TORCH_DTYPE_NAME = {
+    torch.bfloat16: "bfloat16",
+    torch.float16: "float16",
+    torch.float32: "float32",
+}
+
+
+def iter_quantization_configs(obj, path=()):
+    if not isinstance(obj, dict):
+        return
+
+    quant_config = obj.get("quantization_config")
+    if isinstance(quant_config, dict):
+        yield path + ("quantization_config",), quant_config
+
+    for key, value in obj.items():
+        if isinstance(value, dict):
+            yield from iter_quantization_configs(value, path + (key,))
+
+
+def load_compressed_tensors_weight_config(load_path, config_file=None):
+    config_path = Path(config_file).resolve() if config_file is not None else Path(load_path) / "config.json"
+    if not config_path.exists():
+        LOGGER.warning(
+            "compressed-tensors config not found at %s; using fallback Kimi INT4 quantization args.",
+            config_path,
+        )
+        return None, None
+
+    with config_path.open("r", encoding="utf-8") as f:
+        config = json.load(f)
+
+    for path, quant_config in iter_quantization_configs(config):
+        if quant_config.get("quant_method") != "compressed-tensors":
+            continue
+        if quant_config.get("format") != "pack-quantized":
+            continue
+        for group in quant_config.get("config_groups", {}).values():
+            weights = group.get("weights") if isinstance(group, dict) else None
+            if isinstance(weights, dict):
+                return weights, f"{config_path}:{'.'.join(path)}"
+
+    LOGGER.warning(
+        "No pack-quantized compressed-tensors quantization_config found in %s; "
+        "using fallback Kimi INT4 quantization args.",
+        config_path,
+    )
+    return None, None
+
+
+def build_quantization_scheme(
+    load_path,
+    config_file=None,
+    ignore_config_quantization=False,
+    num_bits=4,
+    quant_strategy="group",
+    group_size=32,
+    symmetric=True,
+    dynamic=False,
+):
+    from compressed_tensors.quantization import QuantizationArgs, QuantizationScheme
+
+    weight_config = None
+    config_path = None
+    if not ignore_config_quantization:
+        weight_config, config_path = load_compressed_tensors_weight_config(load_path, config_file)
+
+    if weight_config is None:
+        weight_config = {
+            "num_bits": num_bits,
+            "type": "int",
+            "strategy": quant_strategy,
+            "group_size": group_size,
+            "symmetric": symmetric,
+            "dynamic": dynamic,
+        }
+    else:
+        LOGGER.info("Using compressed-tensors weight config from %s.", config_path)
+
+    return QuantizationScheme(
+        targets=["Linear"],
+        weights=QuantizationArgs(
+            num_bits=weight_config.get("num_bits", num_bits),
+            type=weight_config.get("type", "int"),
+            strategy=weight_config.get("strategy", quant_strategy),
+            group_size=weight_config.get("group_size", group_size),
+            symmetric=weight_config.get("symmetric", symmetric),
+            dynamic=weight_config.get("dynamic", dynamic),
+        ),
+    )
+
+
+def get_packed_weight_keys(
+    weight_key,
+    packed_key=WEIGHT_PACKED_KEY,
+    scale_key=WEIGHT_SCALE_KEY,
+    shape_key=WEIGHT_SHAPE_KEY,
+    zero_point_key=WEIGHT_ZERO_POINT_KEY,
+    g_idx_key=WEIGHT_G_IDX_KEY,
+):
+    base = weight_key[: -len(".weight")] if weight_key.endswith(".weight") else weight_key
+    return [
+        f"{base}.{packed_key}",
+        f"{base}.{scale_key}",
+        f"{base}.{shape_key}",
+        f"{base}.{zero_point_key}",
+        f"{base}.{g_idx_key}",
+    ]
+
+
+def dequantize_int4_packed(
+    weight_packed,
+    weight_scale,
+    weight_shape,
+    quantization_scheme,
+    tensor_name,
+    weight_zero_point=None,
+    weight_g_idx=None,
+):
+    from compressed_tensors import PackedQuantizationCompressor
+
+    if weight_packed.dtype != torch.int32:
+        raise TypeError(
+            f"{tensor_name}.weight_packed must be torch.int32 for compressed-tensors "
+            f"pack-quantized decompression, got {weight_packed.dtype}"
+        )
+
+    compressed_tensors = {
+        "weight_packed": weight_packed,
+        "weight_scale": weight_scale,
+        "weight_shape": weight_shape,
+    }
+    if weight_zero_point is not None:
+        compressed_tensors["weight_zero_point"] = weight_zero_point
+    if weight_g_idx is not None:
+        compressed_tensors["weight_g_idx"] = weight_g_idx
+
+    decompressed = PackedQuantizationCompressor.decompress(
+        compressed_tensors,
+        scheme=quantization_scheme,
+    )
+    return decompressed["weight"]
+
+
+def dequantize_state_dict(
+    state_dict,
+    load_path,
+    output_dtype=torch.bfloat16,
+    config_file=None,
+    ignore_config_quantization=False,
+    packed_key=WEIGHT_PACKED_KEY,
+    scale_key=WEIGHT_SCALE_KEY,
+    shape_key=WEIGHT_SHAPE_KEY,
+    zero_point_key=WEIGHT_ZERO_POINT_KEY,
+    g_idx_key=WEIGHT_G_IDX_KEY,
+    num_bits=4,
+    quant_strategy="group",
+    group_size=32,
+    symmetric=True,
+    dynamic=False,
+    target_weight_keys=None,
+):
+    packed_suffix = f".{packed_key}"
+    if target_weight_keys is None:
+        packed_keys = sorted(key for key in list(state_dict) if key.endswith(packed_suffix))
+    else:
+        packed_keys = []
+        missing_packed_keys = []
+        for weight_key in sorted(target_weight_keys):
+            if not weight_key.endswith(".weight"):
+                raise ValueError(f"Targeted dequant weight key must end with .weight, got: {weight_key}")
+            key = f"{weight_key[:-len('.weight')]}.{packed_key}"
+            if key in state_dict:
+                packed_keys.append(key)
+            else:
+                missing_packed_keys.append(key)
+        if missing_packed_keys:
+            preview = ", ".join(missing_packed_keys[:5])
+            suffix = "" if len(missing_packed_keys) <= 5 else f", ... ({len(missing_packed_keys)} total)"
+            raise KeyError(
+                "Targeted compressed-tensors dequant requested packed tensor(s) that were not loaded: "
+                f"{preview}{suffix}"
+            )
+    if not packed_keys:
+        return 0
+
+    quantization_scheme = build_quantization_scheme(
+        load_path,
+        config_file=config_file,
+        ignore_config_quantization=ignore_config_quantization,
+        num_bits=num_bits,
+        quant_strategy=quant_strategy,
+        group_size=group_size,
+        symmetric=symmetric,
+        dynamic=dynamic,
+    )
+
+    converted = 0
+    total = len(packed_keys)
+    progress_interval = max(1, total // 20)
+    LOGGER.info("Dequantizing %d target compressed-tensors packed INT4 weight(s).", total)
+    for key in packed_keys:
+        base = key[: -len(packed_suffix)]
+        out_key = f"{base}.weight"
+        if out_key in state_dict:
+            raise KeyError(f"Refusing to overwrite existing tensor while dequantizing: {out_key}")
+
+        scale_tensor_key = f"{base}.{scale_key}"
+        shape_tensor_key = f"{base}.{shape_key}"
+        missing = [name for name in (scale_tensor_key, shape_tensor_key) if name not in state_dict]
+        if missing:
+            raise KeyError(f"{key} is missing compressed-tensors companion tensor(s): {missing}")
+
+        tensor = dequantize_int4_packed(
+            state_dict[key],
+            state_dict[scale_tensor_key],
+            state_dict[shape_tensor_key],
+            quantization_scheme,
+            base,
+            weight_zero_point=state_dict.get(f"{base}.{zero_point_key}"),
+            weight_g_idx=state_dict.get(f"{base}.{g_idx_key}"),
+        )
+        state_dict[out_key] = tensor.to(output_dtype).contiguous()
+
+        for remove_key in (
+            key,
+            scale_tensor_key,
+            shape_tensor_key,
+            f"{base}.{zero_point_key}",
+            f"{base}.{g_idx_key}",
+        ):
+            state_dict.pop(remove_key, None)
+        converted += 1
+        if converted == total or converted % progress_interval == 0:
+            LOGGER.info("Dequantized %d/%d compressed-tensors packed INT4 weight(s).", converted, total)
+
+    return converted