[Feature] Add GLM-5.3-Flash F6 core: text model + MTP + compose model

examples/v1/config/sft_glm53.py

@@ -0,0 +1,141 @@
+import os
+
+from xtuner.v1.config import AdamWConfig, FSDPConfig, LRConfig, MuonConfig
+from xtuner.v1.datasets import OpenaiTokenizeFunctionConfig
+from xtuner.v1.datasets.config import DataloaderConfig, DatasetConfig
+from xtuner.v1.float8.config import Float8Config, ScalingGranularity
+from xtuner.v1.loss import CELossConfig
+from xtuner.v1.model import get_model_config_from_hf
+from xtuner.v1.train import TrainerConfig
+from xtuner.v1.train.trainer import LoadCheckpointConfig
+
+
+def _get_bool_env(name: str, default: bool = False) -> bool:
+    return os.environ.get(name, "1" if default else "0").lower() in ("1", "true", "yes", "on")
+
+
+def _get_dispatcher():
+    dispatcher = os.environ.get("DISPATCHER", "all2all").lower()
+    if dispatcher in ("", "none", "null"):
+        return None
+    return dispatcher
+
+
+def _get_float8_config() -> Float8Config | None:
+    if not _get_bool_env("FP8", False):
+        return None
+    return Float8Config(
+        scaling_granularity_gemm=ScalingGranularity.TILEWISE,
+        scaling_granularity_grouped_gemm=ScalingGranularity.TILEWISE,
+    )
+
+
+GLM5_3_MODEL_PATH = os.environ["GLM5_3_MODEL_PATH"]
+
+work_dir = os.environ.get("WORK_DIR", "work_dirs/glm53_sft")
+# On single-node 8-GPU SFT, EP=8 leaves FSDP size at 1 and replicates non-expert params.
+ep_size = int(os.environ.get("EP_SIZE", "1"))
+intra_layer_micro_batch = int(os.environ.get("INTRA_LAYER_MICRO_BATCH", "1"))
+global_batch_size = int(os.environ.get("GLOBAL_BATCH_SIZE", os.environ.get("WORLD_SIZE", "8")))
+sample_max_length = int(os.environ.get("SAMPLE_MAX_LENGTH", "4096"))
+pack_max_length = int(os.environ.get("PACK_MAX_LENGTH", "16384"))
+total_step = int(os.environ.get("TOTAL_STEP", "10"))
+
+loss_cfg = CELossConfig(
+    mode=os.environ.get("LOSS_MODE", "chunk"),
+    chunk_size=int(os.environ.get("LOSS_CHUNK_SIZE", "1024")),
+)
+
+# get_model_config_from_hf dispatches Glm5NextConfig.model_type=="glm5_next" to
+# Glm53TextMoEConfig.from_hf, which reads the nested text_config directly -- this is the
+# text-only SFT path (design doc F6 point 2); VL training goes through Glm53BaseConfig instead.
+model_cfg = get_model_config_from_hf(GLM5_3_MODEL_PATH)
+model_cfg.dispatcher = _get_dispatcher()
+model_cfg.ep_size = ep_size
+model_cfg.compile_cfg = _get_bool_env("MODEL_COMPILE", False)
+model_cfg.float8_cfg = _get_float8_config()
+model_cfg.lm_loss_cfg = loss_cfg
+model_cfg.attention.sparse_mla_backend = os.environ.get("SPARSE_MLA_BACKEND", "flash_mla_cudnn").strip().lower()
+if "INDEXER_BACKEND" in os.environ:
+    model_cfg.attention.indexer_backend = os.environ["INDEXER_BACKEND"].strip().lower()
+
+cache_dir = os.path.join(work_dir, "jsonl_cache")
+cache_tag = os.environ.get("CACHE_TAG", f"glm53_{sample_max_length}")
+dataset_type = os.environ.get("DATASET_TYPE", "alpaca").lower()
+dataset_path_envs = {
+    "alpaca": "ALPACA_PATH",
+    "alpaca_long": "ALPACA_LONG_PATH",
+}
+if dataset_type not in dataset_path_envs:
+    raise ValueError(f"Unsupported DATASET_TYPE={dataset_type!r}. Use alpaca or alpaca_long.")
+dataset_config = [
+    {
+        "dataset": DatasetConfig(
+            name=dataset_type,
+            anno_path=os.environ[dataset_path_envs[dataset_type]],
+            sample_ratio=float(os.environ.get("DATASET_SAMPLE_RATIO", "1.0")),
+            cache_dir=cache_dir,
+            cache_tag=f"{cache_tag}_{dataset_type}",
+        ),
+        "tokenize_fn": OpenaiTokenizeFunctionConfig(
+            chat_template="glm5.3",
+            max_length=sample_max_length,
+        ),
+    }
+]
+
+dataloader_config = DataloaderConfig(
+    dataset_config_list=dataset_config,
+    pack_level=os.environ.get("PACK_LEVEL", "soft"),
+    pack_max_length=pack_max_length,
+    pack_chunk_size=int(os.environ.get("PACK_CHUNK_SIZE", "10000")),
+    pack_workers=int(os.environ.get("PACK_WORKERS", "4")),
+    global_pack=_get_bool_env("GLOBAL_PACK", True),
+    group_by_length=_get_bool_env("GROUP_BY_LENGTH", True),
+    num_workers=int(os.environ.get("DATALOADER_NUM_WORKERS", "4")),
+)
+
+lr = float(os.environ.get("LR", "1e-6"))
+optimizer = os.environ.get("OPTIMIZER", "adamw").lower()
+if optimizer == "muon":
+    optim_cfg = MuonConfig(lr=lr)
+elif optimizer == "adamw":
+    optim_cfg = AdamWConfig(
+        lr=lr,
+        foreach=_get_bool_env("ADAMW_FOREACH", False),
+        swap_optimizer=_get_bool_env("SWAP_OPTIMIZER", False),
+    )
+else:
+    raise ValueError(f"Unsupported OPTIMIZER={optimizer!r}. Use adamw or muon.")
+lr_cfg = LRConfig(lr_type=os.environ.get("LR_TYPE", "cosine"), warmup_ratio=float(os.environ.get("WARMUP_RATIO", "0")))
+fsdp_cfg = FSDPConfig(
+    cpu_offload=_get_bool_env("CPU_OFFLOAD", False),
+    ep_size=ep_size,
+    torch_compile=_get_bool_env("TORCH_COMPILE", False),
+)
+
+trainer = TrainerConfig(
+    model_cfg=model_cfg,
+    load_from=GLM5_3_MODEL_PATH,
+    tokenizer_path=GLM5_3_MODEL_PATH,
+    strict_load=_get_bool_env("STRICT_LOAD", True),
+    optim_cfg=optim_cfg,
+    dataloader_cfg=dataloader_config,
+    lr_cfg=lr_cfg,
+    loss_cfg=loss_cfg,
+    fsdp_cfg=fsdp_cfg,
+    global_batch_size=global_batch_size,
+    total_step=total_step,
+    intra_layer_micro_batch=intra_layer_micro_batch,
+    sp_size=int(os.environ.get("SP_SIZE", "1")),
+    load_checkpoint_cfg=LoadCheckpointConfig(checkpoint_path=os.environ.get("LOAD_CHECKPOINT_PATH")),
+    checkpoint_interval=int(os.environ.get("CHECKPOINT_INTERVAL", "200")),
+    checkpoint_maxkeep=int(os.environ.get("CHECKPOINT_MAX_KEEP", "3")),
+    hf_interval=int(os.environ.get("HF_INTERVAL", "200")),
+    hf_max_keep=int(os.environ.get("HF_MAX_KEEP", "3")),
+    work_dir=work_dir,
+    profile_memory=_get_bool_env("PROFILE_MEMORY", False),
+    profile_time=_get_bool_env("PROFILE_TIME", False),
+    profile_step=[int(x) for x in os.environ.get("PROFILE_STEP", "2,3").split(",") if x],
+    debug_skip_save=_get_bool_env("DEBUG_SKIP_SAVE", False),
+)

examples/v1/config/sft_glm53_vl.py

@@ -0,0 +1,156 @@
+"""GLM-5.3-Flash VL SFT config (image/video), the multimodal counterpart of sft_glm53.py.
+
+The only structural differences from the text-only config are the three that make it a VLM:
+``Glm53BaseConfig.from_hf`` (vision tower + projector + language model, instead of the text half
+alone), ``Glm53VLTokenizeFnConfig`` over a ``VLMJsonlDataset``, and ``glm53_vl_sft_collator``,
+which needs ``image_token_id``/``merge_unit`` bound through ``collator_kwargs`` to re-check its
+placeholder<->patch-count invariant after pack-level truncation.
+"""
+
+import os
+
+from xtuner.v1.config import AdamWConfig, FSDPConfig, LRConfig, MuonConfig
+from xtuner.v1.datasets import Glm53VLTokenizeFnConfig
+from xtuner.v1.datasets.config import DataloaderConfig, DatasetConfig
+from xtuner.v1.float8.config import Float8Config, ScalingGranularity
+from xtuner.v1.loss import CELossConfig
+from xtuner.v1.model.compose.glm53 import Glm53BaseConfig
+from xtuner.v1.train import TrainerConfig
+from xtuner.v1.train.trainer import LoadCheckpointConfig
+
+
+def _get_bool_env(name: str, default: bool = False) -> bool:
+    return os.environ.get(name, "1" if default else "0").lower() in ("1", "true", "yes", "on")
+
+
+def _get_dispatcher():
+    dispatcher = os.environ.get("DISPATCHER", "all2all").lower()
+    if dispatcher in ("", "none", "null"):
+        return None
+    return dispatcher
+
+
+def _get_float8_config() -> Float8Config | None:
+    if not _get_bool_env("FP8", False):
+        return None
+    return Float8Config(
+        scaling_granularity_gemm=ScalingGranularity.TILEWISE,
+        scaling_granularity_grouped_gemm=ScalingGranularity.TILEWISE,
+    )
+
+
+GLM5_3_MODEL_PATH = os.environ["GLM5_3_MODEL_PATH"]
+DATA_PATH = os.environ["VL_DATA_PATH"]
+MEDIA_ROOT = os.environ.get("VL_MEDIA_ROOT", DATA_PATH)
+
+work_dir = os.environ.get("WORK_DIR", "work_dirs/glm53_vl_sft")
+ep_size = int(os.environ.get("EP_SIZE", "1"))
+sp_size = int(os.environ.get("SP_SIZE", "1"))
+intra_layer_micro_batch = int(os.environ.get("INTRA_LAYER_MICRO_BATCH", "1"))
+global_batch_size = int(os.environ.get("GLOBAL_BATCH_SIZE", os.environ.get("WORLD_SIZE", "8")))
+sample_max_length = int(os.environ.get("SAMPLE_MAX_LENGTH", "4096"))
+pack_max_length = int(os.environ.get("PACK_MAX_LENGTH", "16384"))
+total_step = int(os.environ.get("TOTAL_STEP", "10"))
+
+loss_cfg = CELossConfig(
+    mode=os.environ.get("LOSS_MODE", "chunk"),
+    chunk_size=int(os.environ.get("LOSS_CHUNK_SIZE", "1024")),
+)
+
+model_cfg = Glm53BaseConfig.from_hf(GLM5_3_MODEL_PATH)
+model_cfg.text_config.dispatcher = _get_dispatcher()
+model_cfg.text_config.ep_size = ep_size
+model_cfg.text_config.lm_loss_cfg = loss_cfg
+model_cfg.text_config.attention.sparse_mla_backend = (
+    os.environ.get("SPARSE_MLA_BACKEND", "flash_mla_cudnn").strip().lower()
+)
+if "INDEXER_BACKEND" in os.environ:
+    model_cfg.text_config.attention.indexer_backend = os.environ["INDEXER_BACKEND"].strip().lower()
+# FP8 quantizes the language tower's large projections only; the vision half stays bf16, which
+# is also how the published checkpoint stores it.
+model_cfg.text_config.float8_cfg = _get_float8_config()
+model_cfg.vision_config.attn_impl = os.environ.get("VISION_ATTN_IMPL", "flash_attention")
+model_cfg.compile_cfg = _get_bool_env("MODEL_COMPILE", False)
+model_cfg.text_config.compile_cfg = _get_bool_env("MODEL_COMPILE", False)
+
+cache_dir = os.path.join(work_dir, "jsonl_cache")
+dataset_config = [
+    {
+        "dataset": DatasetConfig(
+            name="glm53_vl",
+            anno_path=DATA_PATH,
+            class_name="VLMJsonlDataset",
+            media_root=MEDIA_ROOT,
+            sample_ratio=float(os.environ.get("DATASET_SAMPLE_RATIO", "1.0")),
+            cache_dir=cache_dir,
+            cache_tag=os.environ.get("CACHE_TAG", f"glm53_vl_{sample_max_length}"),
+        ),
+        "tokenize_fn": Glm53VLTokenizeFnConfig(
+            processor_path=GLM5_3_MODEL_PATH,
+            max_length=sample_max_length,
+            max_pixels=int(os.environ["VL_MAX_PIXELS"]) if "VL_MAX_PIXELS" in os.environ else None,
+            min_pixels=int(os.environ["VL_MIN_PIXELS"]) if "VL_MIN_PIXELS" in os.environ else None,
+        ),
+    },
+]
+
+dataloader_config = DataloaderConfig(
+    dataset_config_list=dataset_config,
+    collator="glm53_vl_sft_collator",
+    collator_kwargs={
+        "image_token_id": model_cfg.image_token_id,
+        "merge_unit": model_cfg.vision_config.spatial_merge_size**2,
+    },
+    pack_level=os.environ.get("PACK_LEVEL", "soft"),
+    pack_max_length=pack_max_length,
+    pack_chunk_size=int(os.environ.get("PACK_CHUNK_SIZE", "10000")),
+    pack_workers=int(os.environ.get("PACK_WORKERS", "4")),
+    global_pack=_get_bool_env("GLOBAL_PACK", True),
+    group_by_length=_get_bool_env("GROUP_BY_LENGTH", True),
+    num_workers=int(os.environ.get("DATALOADER_NUM_WORKERS", "4")),
+)
+
+lr = float(os.environ.get("LR", "1e-6"))
+optimizer = os.environ.get("OPTIMIZER", "adamw").lower()
+if optimizer == "muon":
+    optim_cfg = MuonConfig(lr=lr)
+elif optimizer == "adamw":
+    optim_cfg = AdamWConfig(
+        lr=lr,
+        foreach=_get_bool_env("ADAMW_FOREACH", False),
+        swap_optimizer=_get_bool_env("SWAP_OPTIMIZER", False),
+    )
+else:
+    raise ValueError(f"Unsupported OPTIMIZER={optimizer!r}. Use adamw or muon.")
+lr_cfg = LRConfig(lr_type=os.environ.get("LR_TYPE", "cosine"), warmup_ratio=float(os.environ.get("WARMUP_RATIO", "0")))
+fsdp_cfg = FSDPConfig(
+    cpu_offload=_get_bool_env("CPU_OFFLOAD", False),
+    ep_size=ep_size,
+    torch_compile=_get_bool_env("TORCH_COMPILE", False),
+)
+
+trainer = TrainerConfig(
+    model_cfg=model_cfg,
+    load_from=GLM5_3_MODEL_PATH,
+    tokenizer_path=GLM5_3_MODEL_PATH,
+    strict_load=_get_bool_env("STRICT_LOAD", False),
+    optim_cfg=optim_cfg,
+    dataloader_cfg=dataloader_config,
+    lr_cfg=lr_cfg,
+    loss_cfg=loss_cfg,
+    fsdp_cfg=fsdp_cfg,
+    global_batch_size=global_batch_size,
+    total_step=total_step,
+    intra_layer_micro_batch=intra_layer_micro_batch,
+    sp_size=sp_size,
+    load_checkpoint_cfg=LoadCheckpointConfig(checkpoint_path=os.environ.get("LOAD_CHECKPOINT_PATH")),
+    checkpoint_interval=int(os.environ.get("CHECKPOINT_INTERVAL", "200")),
+    checkpoint_maxkeep=int(os.environ.get("CHECKPOINT_MAX_KEEP", "3")),
+    hf_interval=int(os.environ.get("HF_INTERVAL", "200")),
+    hf_max_keep=int(os.environ.get("HF_MAX_KEEP", "3")),
+    work_dir=work_dir,
+    profile_memory=_get_bool_env("PROFILE_MEMORY", False),
+    profile_time=_get_bool_env("PROFILE_TIME", False),
+    profile_step=[int(x) for x in os.environ.get("PROFILE_STEP", "2,3").split(",") if x],
+    debug_skip_save=_get_bool_env("DEBUG_SKIP_SAVE", False),
+)

sft_glm53_tiny.sh

@@ -0,0 +1,80 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+SCRIPT_PATH="${SCRIPT_DIR}/$(basename "$0")"
+cd "${SCRIPT_DIR}"
+
+MODEL_PATH_DEFAULT="/mnt/shared-storage-user/zhaopenghao/model/GLM-5.3-Flash-25B"
+ALPACA_PATH_DEFAULT="/mnt/shared-storage-user/llmrazor-share/data/alpaca"
+ALPACA_LONG_PATH_DEFAULT="/mnt/shared-storage-user/zhaopenghao/slime0701/datasets/alpaca_long_glm5_t6144_m8192.jsonl"
+
+# TinyGLM53 profile: EP4 / flash_mla_cudnn DSA / 16K pack; single-node by default.
+export CONDA_ENV="${CONDA_ENV:-pt29_glm2}"
+export GLM5_3_MODEL_PATH="${GLM5_3_MODEL_PATH:-${MODEL_PATH_DEFAULT}}"
+export ALPACA_PATH="${ALPACA_PATH:-${ALPACA_PATH_DEFAULT}}"
+export ALPACA_LONG_PATH="${ALPACA_LONG_PATH:-${ALPACA_LONG_PATH_DEFAULT}}"
+export CONFIG_PATH="${CONFIG_PATH:-examples/v1/config/sft_glm53.py}"
+export WORK_DIR="${WORK_DIR:-work_dirs/sft_glm53/tiny}"
+export PYTHONPATH="./${PYTHONPATH:+:${PYTHONPATH}}"
+export NNODES="${NNODES:-1}"
+export NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
+export NODE_RANK="${NODE_RANK:-0}"
+export MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
+export MASTER_PORT="${MASTER_PORT:-29671}"
+
+export EP_SIZE="${EP_SIZE:-4}"
+export SP_SIZE="${SP_SIZE:-1}"
+export DISPATCHER="${DISPATCHER:-all2all}"
+export SPARSE_MLA_BACKEND="${SPARSE_MLA_BACKEND:-flash_mla_cudnn}"
+
+export DATASET_TYPE="${DATASET_TYPE:-alpaca}"
+export DATASET_SAMPLE_RATIO="${DATASET_SAMPLE_RATIO:-1.0}"
+export SAMPLE_MAX_LENGTH="${SAMPLE_MAX_LENGTH:-4096}"
+export PACK_MAX_LENGTH="${PACK_MAX_LENGTH:-16384}"
+# Keep one packed sample per data-parallel rank when scaling out with SP.
+export GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-$((NNODES * NPROC_PER_NODE / SP_SIZE))}"
+export INTRA_LAYER_MICRO_BATCH="${INTRA_LAYER_MICRO_BATCH:-1}"
+export TOTAL_STEP="${TOTAL_STEP:-20}"
+
+export FP8="${FP8:-0}"
+export MODEL_COMPILE="${MODEL_COMPILE:-1}"
+export TORCH_COMPILE="${TORCH_COMPILE:-1}"
+export LOSS_CHUNK_SIZE="${LOSS_CHUNK_SIZE:-2048}"
+export SWAP_OPTIMIZER="${SWAP_OPTIMIZER:-0}"
+# GLM-5.3-Flash's checkpoint is a VL checkpoint (vision + projector + language_model); this
+# script loads only Glm53TextMoEConfig (text half), so `model.visual.*` keys are always
+# "unexpected" under strict loading -- that's expected for a text-only load, not a bug. Unlike
+# GLM-5.2 (text-only checkpoint, where STRICT_LOAD=1 is correct), default this to 0.
+export STRICT_LOAD="${STRICT_LOAD:-0}"
+export DEBUG_SKIP_SAVE="${DEBUG_SKIP_SAVE:-1}"
+
+export XTUNER_ACTIVATION_OFFLOAD="${XTUNER_ACTIVATION_OFFLOAD:-1}"
+export XTUNER_GC_ENABLE="${XTUNER_GC_ENABLE:-1}"
+export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
+
+export PROFILE_TIME="${PROFILE_TIME:-0}"
+export PROFILE_MEMORY="${PROFILE_MEMORY:-0}"
+export PROFILE_STEP="${PROFILE_STEP:-8}"
+
+source "/mnt/shared-storage-user/zhaopenghao/miniconda3/etc/profile.d/conda.sh"
+conda activate "${CONDA_ENV}"
+
+mkdir -p "${WORK_DIR}"
+if [[ "${NODE_RANK}" == "0" ]]; then
+  cp -f "${SCRIPT_PATH}" "${WORK_DIR}/$(basename "${SCRIPT_PATH}")"
+  cp -f "${CONFIG_PATH}" "${WORK_DIR}/$(basename "${CONFIG_PATH}")"
+fi
+
+current_time=$(date "+%m%d%H%M%S")
+env | grep -E '^(CONDA_ENV|GLM5_3_MODEL_PATH|ALPACA_PATH|ALPACA_LONG_PATH|DATASET_TYPE|DATASET_SAMPLE_RATIO|CONFIG_PATH|WORK_DIR|NNODES|NPROC_PER_NODE|NODE_RANK|MASTER_ADDR|MASTER_PORT|SAMPLE_MAX_LENGTH|PACK_MAX_LENGTH|GLOBAL_BATCH_SIZE|INTRA_LAYER_MICRO_BATCH|TOTAL_STEP|EP_SIZE|SP_SIZE|DISPATCHER|SPARSE_MLA_BACKEND|FP8|MODEL_COMPILE|TORCH_COMPILE|LOSS_CHUNK_SIZE|SWAP_OPTIMIZER|STRICT_LOAD|DEBUG_SKIP_SAVE|PYTORCH_CUDA_ALLOC_CONF|PROFILE_[A-Z0-9_]+)=' | sort
+
+torchrun \
+  --nnodes "${NNODES}" \
+  --nproc-per-node "${NPROC_PER_NODE}" \
+  --node-rank "${NODE_RANK}" \
+  --master-addr "${MASTER_ADDR}" \
+  --master-port "${MASTER_PORT}" \
+  xtuner/v1/train/cli/sft.py \
+  --config "${CONFIG_PATH}" \
+  2>&1 | tee -a "${WORK_DIR}/training_log_${current_time}_${NODE_RANK}.txt"