マルチノードでAxolotl を使ってQwen3-235B-A22B-Thinking-2507にHARIする方法について

はじめに

本稿は Axolotl をマルチノードで回しながら、Qwen3-235B-A22B-Thinking-2507 に HARI を適用する実運用手順をまとめたものです。HARI は「LoRA を有効にした学習損失に、LoRA を一時無効化したベース参照出力との正則化を加える」軽量 SFT で、ここでは KV/HS 自己蒸留として実装します。Axolotl では trainer_cls で カスタムトレーナーを差し替え、挙動は YAML と 環境変数で切り替えます。つまりコードは一度置けば、実験条件は YAML＋環境変数の編集だけで管理できます。

1. 全体像

まず カスタムトレーナーを Axolotl から参照できる場所に置きます。次に YAML に trainer_cls を指定し、LoRA/データ/分散の設定を記述します。最後に 環境変数で KV 蒸留のノブ（α・層・トークン間引き・教師の有無）を与え、sbatch で起動します。以降は W&B / STDOUT に loss_ce と loss_kvsd が並んで出るので、HARI の効き目をそのまま可視化できます。

1.1. カスタムトレーナー（KV/HS 自己蒸留）

すでに用意済みの kv_distill_trainer.py を Axolotl から import できるパスに配置します（例：/home/Competition2025/P10/P10U001/work/axolotl/kv_distill_trainer.py）。このトレーナーは次を行います。

• 学生パスは通常どおり LoRA 有効で forward。
• 教師パスは既定で 共有教師（同一モデルの LoRA 一時無効化＋eval）を使用。必要なら環境変数で別教師を自動ロード。
• まず k_proj / v_proj の出力をフックして KV-MSE を計算。もし取れなければ Hidden States 蒸留にフォールバック。
• 各 step で loss_ce / loss_kvsd / kvsd_alpha を self.log に出力（W&B でも可視化）。
• Axolotl 特有の追加引数を吸収して HF Trainer へ安全に橋渡し。

1.2. YAML（Axolotl 設定）

trainer_cls に先ほどのクラスを指定します。LoRA は Qwen3-235B の self_attn の q/k/v/o を対象に、長文志向のため sequence_len=16384、sample_packing を有効化。分散は FSDP を選択しています（必要に応じて Deepspeed に差し替え可）。

# /home/Competition2025/P10/P10U001/work/axolotl/axolotl_qwen3_235b_fsdp.yaml
base_model: /nvme34/Qwen/Qwen3-235B-A22B-Thinking-2507
model_type: AutoModelForCausalLM
tokenizer_type: AutoTokenizer
trainer_cls: kv_distill_trainer.KVSDTrainer

# 省略: Liger 最適化, 4bit 読み込み, データ, LoRA(q/k/v/o), 16k 設定, FSDP…

要点は次の三つです。

1. trainer_cls: kv_distill_trainer.KVSDTrainer で蒸留対応 Trainer を有効化。
2. LoRA の lora_target_modules を self_attn の q/k/v/o に限定し、容量と効果のバランスを取る。
3. 長コンテキストのバッチ作成を安定化するため sequence_len=16384 + sample_packing: true をセット。

1.3. 環境変数（HARI のノブはここで回す）

実験の切り替えは 環境変数で即時反映されます。sbatch スクリプト内で設定しておくと、YAMLを汚さずに条件を変えられます。

# sbatch 内（抜粋）
export PYTHONPATH="/home/Competition2025/P10/P10U001/work/axolotl:$PYTHONPATH"

# --- HARI (KV/HS 自己蒸留) ノブ ---
export AXO_KVSD_ALPHA=0.9           # ← HARI の要。0.5〜0.9が実戦域、まずは 0.9 を推奨
export AXO_KVSD_LAYERS=all          # all / every2 / last4
export AXO_KVSD_TOKEN_STRIDE=4      # KV/HS 抽出のトークン間引きで VRAM/速度を調整
export AXO_KVSD_SHARE_STUDENT=true  # 共有教師（LoRA 無効化）を使う
export AXO_KVSD_HS_FALLBACK=true    # KV が取れない実装では HS に自動フォールバック

# 別教師を使う場合（任意）
# export AXO_KVSD_TEACHER_MODEL="Qwen/Qwen2.5-7B-Instruct"
# export AXO_KVSD_TEACHER_4BIT=true

HARI のキモは AXO_KVSD_ALPHA です。ベース保持を最優先するなら 0.3–0.6、短期で底上げを狙うなら 0.9 をまず試す、という運用が分かりやすいです。AXO_KVSD_TOKEN_STRIDE は 2–8 程度で VRAM と速度をトレードオフ可能。

1.4. ジョブスクリプトと実行

srun/torchrun で ノード間 DDP を張り、各ノードは --nproc_per_node=${SLURM_GPUS_PER_NODE} でフル活用します。W&B を有効にしていればブラウザから loss_kvsd の推移を確認できます。

ログには学習の通常損失に加えて loss_kvsd が並びます。loss_ce が下がりつつ loss_kvsd が一定の張り付きで推移すれば、ベース挙動を保ちつつ微調整が入っているサインです。

2. 環境構築

以下のcondaを作成してください。
Condaバージョン: Miniconda 24.7.1
Python: 3.11

# ==== 0) 既存envクリア ====
conda deactivate 2>/dev/null || true
rm -rf "/home/Competition2025/P10/P10U001/.conda/envs/axo-fsdp"
conda env remove -n axo-fsdp -y || true
conda clean -afy || true
pip cache purge || true

# ==== 1) Conda環境 (Pythonは 3.10 のまま) ====
conda create -y -n axo-fsdp python=3.10 cmake ninja pkg-config pip git
eval "$($(conda info --base)/bin/conda shell.bash hook)"
conda activate axo-fsdp

# ==== 2) PyTorch 2.7.1 + cu128 固定 ====
pip install -U --extra-index-url https://download.pytorch.org/whl/cu128 \
  "torch==2.7.1+cu128" "torchvision==0.22.1+cu128" "torchaudio==2.7.1+cu128"

# ==== 3) CUDA Toolkit (NVCC入り) を conda で導入 ====
#  - v12.8 をTorch(=cu128)に合わせて導入
#  - 必要なヘッダ/ライブラリと nvcc を同梱
conda install -y -c nvidia -c conda-forge \
  "cuda-toolkit=12.8" "cuda-nvcc=12.8"

# 推奨: 明示的に環境変数を通す（DeepSpeed のビルドが拾いやすくなる）
export CUDA_HOME="$CONDA_PREFIX"
export CUDA_PATH="$CONDA_PREFIX"
export PATH="$CONDA_PREFIX/bin:$PATH"
export LD_LIBRARY_PATH="$CONDA_PREFIX/lib:$CONDA_PREFIX/lib64:${LD_LIBRARY_PATH:-}"

# （任意）GPUアーキが分かっていれば指定（例: H100なら 90a）
# export TORCH_CUDA_ARCH_LIST="90;90a"

# ==== 4) Axolotl (extrasで flash-attn / deepspeed / vllm を一括) ====
pip install -U packaging setuptools wheel ninja

set -e
if ! pip install -U "axolotl[flash-attn,deepspeed,vllm]==0.12.2"; then
  echo "[WARN] extras 失敗。最小構成へフォールバックします。"
  pip install -U "axolotl==0.12.2"
  # 高速化/DSを個別に試行（Wheel無い/ビルド不可ならスキップ）
  pip install -U "flash-attn>=2.6.0" || true
  pip install -U "deepspeed==0.14.4" || true
fi
set +e

# ==== 5) Axolotl が入れない“最低限の追加”のみ ====
# 5-1) 4bit/QLoRA 用
pip install -U "bitsandbytes==0.47.0" "sentencepiece>=0.2.0" "einops>=0.7.0" \
               "protobuf>=4.25,<6" "pynvml>=11.5.0" "psutil>=5.9" "safetensors>=0.4.2"

# 5-2) vLLM外部API連携（OpenAI互換クライアント）
pip install -U "openai>=1.40.0" "httpx>=0.27.0" "aiohttp>=3.9" "orjson>=3.10" "uvloop>=0.19"

# 5-3) ログ/可視化（任意）
pip install -U "wandb>=0.17.0" "tensorboard>=2.16"

pip uninstall -y torchvision

pip install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@0ee9ee8"

# 整合チェック
pip check || true

# ==== 6) 動作確認 ====
echo "[nvcc]"; nvcc --version || true
python - <<'PY'
import torch, importlib.metadata as md, pkgutil, sys, os
def v(n):
    try: print(f"{n}: {md.version(n)}")
    except:  print(f"{n}: not installed")
print("python:", sys.version.split()[0])
print("CUDA runtime (torch):", torch.version.cuda, "| available:", torch.cuda.is_available(), "| nGPU:", torch.cuda.device_count())
print("CUDA_HOME:", os.getenv("CUDA_HOME"))
for p in ["axolotl","transformers","accelerate","datasets","tokenizers","safetensors","huggingface_hub",
          "peft","trl","bitsandbytes","flash_attn","deepspeed","vllm","openai","wandb"]:
    v(p)
PY

# bnbのCUDA検出（任意）
python -m bitsandbytes | sed -n '1,80p'

3. HARI

3.1. カスタムトレーナー

cat > /home/Competition2025/P10/P10U001/work/axolotl/kv_distill_trainer.py <<'PY'
# kv_distill_trainer.py
# Axolotl `trainer_cls` で指定して使う、KV/HS自己蒸留 Trainer（W&B/STDOUT ログ付き）
# - 既定は「共有教師（LoRA一時無効 + eval）」での自己蒸留
# - 必要なら環境変数で別教師を自動ロード（AXO_KVSD_TEACHER_MODEL）
# - 低VRAM向けにトークン間引き/層選択あり
# - 学習開始時に有効化バナー、最初のバッチでフック数を表示
# - 各ステップで loss_ce / loss_kvsd / kvsd_alpha を self.log に出力（W&B 可視化）
# - 互換修正: eval_data_collator/bench_data_collator や Axolotl 固有の追加 kwargs を吸収して HF Trainer に渡さない

from __future__ import annotations
import os
import re
import inspect
from contextlib import nullcontext
from typing import Dict, Any, Tuple, List, Optional

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from transformers import Trainer

try:
    # ある環境ではBitsAndBytesが無い場合があるため、任意利用
    from transformers import BitsAndBytesConfig, AutoConfig, AutoModelForCausalLM
    _HF_OK = True
except Exception:
    _HF_OK = False
    BitsAndBytesConfig = object  # type: ignore
    AutoConfig = object          # type: ignore
    AutoModelForCausalLM = object  # type: ignore


# ---------------- utils ----------------
def _is_main_process() -> bool:
    try:
        import torch.distributed as dist
        return (not dist.is_available()) or (not dist.is_initialized()) or dist.get_rank() == 0
    except Exception:
        # deepspeed-zero3等でもlocal_rankが-1 or 0のときはTrue扱い
        return True


def _env_flag(name: str, default: bool) -> bool:
    v = os.environ.get(name, None)
    if v is None:
        return default
    return str(v).lower() in ("1", "true", "yes", "y", "on")


# ---------------- KV/HS 取得・損失 ----------------
def _collect_kv_outputs(
    model, inputs: Dict[str, Any], token_stride: int = 1, capture_once_per_module: bool = True
) -> Tuple[Any, Dict[str, List[torch.Tensor]], bool]:
    """
    k_proj / v_proj の線形出力（RoPE前）をフックで取得。
    - gradient checkpoint等でforwardが再入する場合があるため、
      capture_once_per_module=True のときは「同モジュールの最初の呼び出しのみ」採用。
    """
    kv: Dict[str, List[torch.Tensor]] = {"k": [], "v": []}
    handles = []
    seen_k, seen_v = set(), set()

    def make_hook(kind: str, module_name: str):
        def _hook(_m, _inp, out):
            # out: [B, T, D] を想定（多くのQwen/Llama実装で該当）
            if capture_once_per_module:
                bucket = seen_k if kind == "k" else seen_v
                if module_name in bucket:
                    return
                bucket.add(module_name)
            t = out
            if token_stride > 1 and t.ndim >= 3:
                t = t[:, ::token_stride, :]
            kv[kind].append(t)
        return _hook

    for name, module in model.named_modules():
        # Qwen/Llama系は .self_attn.k_proj / v_proj が一般的
        if name.endswith("k_proj") and hasattr(module, "register_forward_hook"):
            handles.append(module.register_forward_hook(make_hook("k", name)))
        elif name.endswith("v_proj") and hasattr(module, "register_forward_hook"):
            handles.append(module.register_forward_hook(make_hook("v", name)))

    outputs = model(**inputs)

    for h in handles:
        try:
            h.remove()
        except Exception:
            pass

    found = len(kv["k"]) > 0 and len(kv["v"]) > 0
    return outputs, kv, found


def _hidden_states_distill(
    student_out, teacher_out, layers_mode: str = "last4", token_stride: int = 1, ref_device: Optional[torch.device] = None
) -> torch.Tensor:
    if not (getattr(student_out, "hidden_states", None) and getattr(teacher_out, "hidden_states", None)):
        dev = ref_device or student_out.logits.device
        return torch.tensor(0.0, device=dev)

    s_hs = student_out.hidden_states
    t_hs = teacher_out.hidden_states

    def _select(seq, mode):
        if mode == "every2":
            return seq[::2]
        if mode == "last4":
            return seq[-4:]
        return seq

    s_sel = _select(s_hs, layers_mode)
    t_sel = _select(t_hs, layers_mode)

    dev = ref_device or student_out.logits.device
    loss = None
    n = 0
    for s, t in zip(s_sel, t_sel):
        if token_stride > 1 and s.ndim >= 3 and t.ndim >= 3:
            s = s[:, ::token_stride, :]
            t = t[:, ::token_stride, :]
        d = min(s.size(-1), t.size(-1))
        mse = F.mse_loss(s[..., :d], t[..., :d])
        loss = mse if loss is None else (loss + mse)
        n += 1
    return (loss / n) if (loss is not None and n > 0) else torch.tensor(0.0, device=dev)


def _kv_mse(
    student_kv: Dict[str, List[torch.Tensor]],
    teacher_kv: Dict[str, List[torch.Tensor]],
    layers_mode: str = "all",
    ref_device: Optional[torch.device] = None,
) -> torch.Tensor:
    def _select(seq, mode):
        if mode == "every2":
            return seq[::2]
        if mode == "last4":
            return seq[-4:]
        return seq

    # どちらかが空なら0
    if not student_kv["k"] or not teacher_kv["k"]:
        dev = ref_device or (student_kv["v"][0].device if student_kv["v"] else torch.device("cpu"))
        return torch.tensor(0.0, device=dev)

    dev = ref_device or student_kv["k"][0].device
    loss = None
    n = 0

    s_k = _select(student_kv["k"], layers_mode)
    s_v = _select(student_kv["v"], layers_mode)
    t_k = _select(teacher_kv["k"], layers_mode)
    t_v = _select(teacher_kv["v"], layers_mode)

    for sk, tk in zip(s_k, t_k):
        d = min(sk.size(-1), tk.size(-1))
        mse = F.mse_loss(sk[..., :d], tk[..., :d])
        loss = mse if loss is None else (loss + mse)
        n += 1

    for sv, tv in zip(s_v, t_v):
        d = min(sv.size(-1), tv.size(-1))
        mse = F.mse_loss(sv[..., :d], tv[..., :d])
        loss = mse if loss is None else (loss + mse)
        n += 1

    return (loss / n) if (loss is not None and n > 0) else torch.tensor(0.0, device=dev)


# ---------------- Trainer ----------------
class KVSDTrainer(Trainer):
    """
    KV/HS 自己蒸留 Trainer for Axolotl (custom trainer_cls).
    - 共有教師: 学生モデルに対して一時的に LoRA/adapter を無効化し、eval() で dropout を止めて教師パスを取る
    - 別教師: AXO_KVSD_TEACHER_MODEL を指定すると別のHFモデルをロードして教師にできる（任意で4bit）
    - 蒸留信号: まず k_proj / v_proj 出力のMSE（KV蒸留）。取れない場合はHS蒸留にフォールバック可能。
    - 互換: Transformers の Trainer は eval_data_collator 等の追加 kwargs を受け付けないため、本クラスで吸収する。
    """

    def __init__(self, *args, **kwargs):
        # ---- 互換修正: ここで Axolotl 由来の追加 kwargs を吸収 ----
        self.eval_data_collator = kwargs.pop("eval_data_collator", None)
        self.bench_data_collator = kwargs.pop("bench_data_collator", None)

        # HF Trainer の __init__ が受け付ける引数だけにフィルタ
        dropped = {}
        try:
            allowed = set(inspect.signature(Trainer.__init__).parameters.keys()) - {"self"}
            filtered = {}
            for k, v in kwargs.items():
                if k in allowed:
                    filtered[k] = v
                else:
                    dropped[k] = v
            kwargs = filtered
        except Exception:
            # 署名取得に失敗したら後段の例外パースに任せる
            pass

        # それでも残る未知キーワードがあれば、TypeError メッセージを解析して順次ドロップ
        try:
            super().__init__(*args, **kwargs)
        except TypeError as e:
            m = re.search(r"unexpected keyword argument '([^']+)'", str(e))
            while m:
                bad = m.group(1)
                dropped[bad] = kwargs.pop(bad, None)
                try:
                    super().__init__(*args, **kwargs)
                    break
                except TypeError as e2:
                    m = re.search(r"unexpected keyword argument '([^']+)'", str(e2))
            else:
                # 予期しない TypeError はそのまま投げる
                raise

        self._dropped_init_kwargs = sorted(dropped.keys())
        if _is_main_process() and dropped:
            print(f"[KVSD] dropped unsupported Trainer kwargs: {self._dropped_init_kwargs}")

        # ---- ハイパラ（環境変数） ----
        self.kv_alpha: float = float(os.environ.get("AXO_KVSD_ALPHA", "0.1"))
        self.kv_layers: str = os.environ.get("AXO_KVSD_LAYERS", "all")  # "all" | "every2" | "last4"
        self.kv_stride: int = int(os.environ.get("AXO_KVSD_TOKEN_STRIDE", "4"))
        self.share_student: bool = _env_flag("AXO_KVSD_SHARE_STUDENT", True)
        self.use_hs_fallback: bool = _env_flag("AXO_KVSD_HS_FALLBACK", True)

        # ---- 別教師（任意） ----
        self.teacher_model = None
        teacher_id = os.environ.get("AXO_KVSD_TEACHER_MODEL", "").strip()
        if teacher_id:
            if not _HF_OK:
                raise RuntimeError("AXO_KVSD_TEACHER_MODEL が指定されましたが、transformersのAuto* importに失敗しています。")
            quant4 = _env_flag("AXO_KVSD_TEACHER_4BIT", True)
            bnb = None
            if quant4:
                try:
                    bnb = BitsAndBytesConfig(
                        load_in_4bit=True, bnb_4bit_quant_type="nf4",
                        bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True
                    )
                except Exception:
                    bnb = None  # bnbが無ければFP16/BF16で落とす
            cfg = AutoConfig.from_pretrained(teacher_id, trust_remote_code=True, token=os.environ.get("HUGGING_FACE_HUB_TOKEN"))
            cfg.use_cache = False  # 学習互換のため
            self.teacher_model = AutoModelForCausalLM.from_pretrained(
                teacher_id,
                config=cfg,
                device_map="auto",
                low_cpu_mem_usage=True,
                trust_remote_code=True,
                quantization_config=bnb,
                token=os.environ.get("HUGGING_FACE_HUB_TOKEN"),
            )
            self.teacher_model.eval()
            for p in self.teacher_model.parameters():
                p.requires_grad = False

        # ---- ログ制御 ----
        self._kvsd_hooks_once = False
        self._kvsd_last_logged_step = -1

        if _is_main_process():
            print(
                f"[KVSD] enabled: alpha={self.kv_alpha} layers={self.kv_layers} "
                f"stride={self.kv_stride} share_student={self.share_student} "
                f"hs_fallback={self.use_hs_fallback} teacher={'external' if self.teacher_model is not None else 'shared'}"
            )

    # 共有教師パス：LoRA/adapterを一時無効にし dropoutを止めてforward（KVのみ取得）
    def _teacher_forward_shared_kv(self, model, inputs: Dict[str, Any]):
        try:
            ctx = model.disable_adapter()
        except AttributeError:
            ctx = nullcontext()
        training = model.training
        with torch.no_grad():
            with ctx:
                model.eval()
                # ★ OOM回避: hidden_statesは要求しない
                t_outputs, t_kv, t_has_kv = _collect_kv_outputs(
                    model, {**inputs, "output_hidden_states": False}, token_stride=self.kv_stride
                )
                model.train(training)
        return t_outputs, t_kv, t_has_kv

    def _maybe_print_hook_stats_once(self, s_kv, t_kv, s_has_kv: bool, t_has_kv: bool):
        if _is_main_process() and (not self._kvsd_hooks_once):
            sk = len(s_kv["k"]); sv = len(s_kv["v"])
            tk = (len(t_kv["k"]) if t_kv else 0); tv = (len(t_kv["v"]) if t_kv else 0)
            print(f"[KVSD] hooks: student(k={sk}, v={sv}) teacher(k={tk}, v={tv})")
            if not (s_has_kv and t_has_kv) and self.use_hs_fallback:
                print("[KVSD] falling back to hidden-states distillation")
            self._kvsd_hooks_once = True

    # num_items_in_batchはTransformers>=4.44系で来ることがあるので受け取っておく
    def compute_loss(self, model, inputs, return_outputs: bool = False, num_items_in_batch: Optional[int] = None, **kwargs):
        # ---- 教師パス（まずはKVのみ、hidden_statesは要求しない）----
        t_outputs = t_kv = None
        t_has_kv = False
        if self.teacher_model is not None:
            with torch.no_grad():
                t_outputs, t_kv, t_has_kv = _collect_kv_outputs(
                    self.teacher_model, {**inputs, "output_hidden_states": False}, token_stride=self.kv_stride
                )
        elif self.share_student:
            t_outputs, t_kv, t_has_kv = self._teacher_forward_shared_kv(model, {**inputs})

        # ---- 学生パス（KVのみ）----
        outputs, s_kv, s_has_kv = _collect_kv_outputs(
            model, {**inputs, "output_hidden_states": False}, token_stride=self.kv_stride
        )
        ce_loss = outputs.loss
        ref_dev = ce_loss.device

        # ---- 蒸留ロス ----
        distill_loss = torch.tensor(0.0, device=ref_dev)

        # ---- 初回のフック状況ログ ----
        self._maybe_print_hook_stats_once(s_kv, t_kv, s_has_kv, t_has_kv)

        if t_outputs is not None:
            if s_has_kv and t_has_kv:
                # 通常はこれでOK（軽い）
                distill_loss = _kv_mse(s_kv, t_kv, layers_mode=self.kv_layers, ref_device=ref_dev)
            elif self.use_hs_fallback:
                # ここだけ重いので、一旦参照を解放してから再forward
                try:
                    del outputs, t_outputs
                    torch.cuda.empty_cache()
                except Exception:
                    pass

                # --- 教師: HS 用に no_grad で再 forward（共有教師はLoRA一時無効）
                if self.teacher_model is not None:
                    with torch.no_grad():
                        t_outputs_hs = self.teacher_model(**{**inputs, "output_hidden_states": True})
                else:
                    try:
                        ctx = model.disable_adapter()
                    except AttributeError:
                        ctx = nullcontext()
                    training = model.training
                    with torch.no_grad():
                        with ctx:
                            model.eval()
                            t_outputs_hs = model(**{**inputs, "output_hidden_states": True})
                            model.train(training)

                # --- 学生: HS 用に勾配ありで再 forward
                outputs_hs = model(**{**inputs, "output_hidden_states": True})

                mode = "last4" if self.kv_layers == "last4" else ("every2" if self.kv_layers == "every2" else "all")
                distill_loss = _hidden_states_distill(
                    outputs_hs, t_outputs_hs, layers_mode=mode, token_stride=self.kv_stride, ref_device=ref_dev
                )
                # 以降の返却 outputs は HS 付きのものに差し替え
                outputs = outputs_hs

        loss = ce_loss + self.kv_alpha * distill_loss

        # ---- ログ出力（各ステップ1回） ----
        try:
            gs = getattr(self.state, "global_step", None)
            if _is_main_process() and (gs is None or gs != self._kvsd_last_logged_step):
                self.log({
                    "loss_ce": float(ce_loss.detach().item()),
                    "loss_kvsd": float(distill_loss.detach().item()),
                    "kvsd_alpha": float(self.kv_alpha),
                })
                self._kvsd_last_logged_step = gs
        except Exception:
            pass

        return (loss, outputs) if return_outputs else loss

    # ---- 追加: 評価時だけ eval_data_collator を使う ----
    def get_eval_dataloader(self, eval_dataset=None):
        eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
        if eval_dataset is None:
            raise ValueError("Trainer: evaluation requires an eval_dataset.")
        # collatorが未指定なら既定実装へ
        if self.eval_data_collator is None:
            return super().get_eval_dataloader(eval_dataset)

        sampler = self._get_eval_sampler(eval_dataset)
        return DataLoader(
            eval_dataset,
            batch_size=self.args.eval_batch_size,
            sampler=sampler,
            collate_fn=self.eval_data_collator,
            drop_last=self.args.dataloader_drop_last,
            num_workers=self.args.dataloader_num_workers,
            pin_memory=self.args.dataloader_pin_memory,
            persistent_workers=self.args.dataloader_persistent_workers,
            timeout=self.args.dataloader_timeout,
            worker_init_fn=self._data_loader_worker_init,
            prefetch_factor=self.args.dataloader_prefetch_factor if self.args.dataloader_prefetch_factor and self.args.dataloader_prefetch_factor > 0 else None,
        )
PY

3.2. yaml

tmp=$(mktemp) && cat > "$tmp" <<'EOF'
base_model: /nvme34/Qwen/Qwen3-235B-A22B-Thinking-2507
model_type: AutoModelForCausalLM
tokenizer_type: AutoTokenizer
trainer_cls: kv_distill_trainer.KVSDTrainer

hub_model_id:
hub_strategy:
push_dataset_to_hub:
hf_use_auth_token: true

# Liger Kernelの設定
plugins:
  - axolotl.integrations.liger.LigerPlugin
  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin

# liger_cross_entropy: true
liger_rope: true
liger_rms_norm: true
liger_glu_activation: true
liger_fused_linear_cross_entropy: true
# cut_cross_entropy: true
liger_layer_norm: true

load_in_8bit: false
load_in_4bit: true
strict: false

chat_template: tokenizer_default

datasets:
  - path: oNo-1/difficult_problem_dataset_v4
    split: "train[:66%]"
    type: chat_template
    field_messages: messages
    roles_to_train: ["assistant"]

dataset_processes: 128

shuffle_merged_datasets: true
dataset_prepared_path: /home/Competition2025/P10/P10U001/data/oNo-1/difficult_problem_dataset_v4
# val_set_size: 0.005
output_dir: /home/Competition2025/P10/P10U001/models/oNo-1/difficult_problem_dataset_v4_kv

sequence_len: 16384
sample_packing: true
eval_sample_packing: false
pad_to_sequence_len: true

adapter: qlora
lora_model_dir:
lora_r: 128
lora_alpha: 256
lora_dropout: 0.05
# lora_target_linear:
lora_target_modules:
  - self_attn.q_proj
  - self_attn.k_proj
  - self_attn.v_proj
  - self_attn.o_proj
  # - down_proj
  # - gate_proj
  # - up_proj
lora_modules_to_save:
  - embed_tokens
  # - lm_head
lora_fan_in_fan_out:
# lora_mlp_kernel: true
# lora_qkv_kernel: true
# lora_o_kernel: true

wandb_project: Qwen3-235B-A22B-Thinking-2507
wandb_entity: llm-competition-2025-wandb-01-
wandb_watch:
wandb_name:
wandb_log_model:

gradient_accumulation_steps: 1
micro_batch_size: 1
num_epochs: 1
optimizer: adamw_torch_4bit
lr_scheduler: cosine
cosine_min_lr_ratio: 0.01
learning_rate: 1e-5
max_grad_norm: 1.0

train_on_inputs: false
group_by_length: false
bf16: true
fp16: false
tf32: false
bfloat16: true

# gradient_checkpointing: true
# gradient_checkpointing_kwargs:
#   use_reentrant: true
early_stopping_patience:
auto_resume_from_checkpoints: true
local_rank:
logging_steps: 1
xformers_attention:
flash_attention: false

save_strategy: steps
save_steps: 100000
save_total_limit: 2
save_only_model: true
# eval_strategy: steps
# eval_steps: 100
# eval_batch_size: 8
# low_cpu_mem_usage: true

warmup_ratio: 0.03
debug:
# deepspeed: /home/Competition2025/P10/P10U001/work/axolotl/deepspeed_configs/zero3_bf16.json
weight_decay: 0.01
fsdp:
fsdp_config:
# metric_for_best_model: eval_loss
# load_best_model_at_end: true

fsdp_version: 1
fsdp:
  - full_shard
  - auto_wrap
fsdp_config:
  fsdp_limit_all_gathers: true
  fsdp_sync_module_states: true
  fsdp_offload_params: false
  fsdp_use_orig_params: false
  fsdp_cpu_ram_efficient_loading: true
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_transformer_layer_cls_to_wrap: Qwen3MoeDecoderLayer
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_activation_checkpointing: true
qlora_sharded_model_loading: true

# DataLoader settings
dataloader_num_workers: 2
dataloader_prefetch_factor: 2
dataloader_pin_memory: true
dataloader_persistent_workers: true

ddp_timeout: 180000000
lora_on_cpu: false
EOF
mv -f "$tmp" /home/Competition2025/P10/P10U001/work/axolotl/axolotl_qwen3_235b_fsdp.yaml

3.3. sbatch

# axolotl_train_qwen3.sh を作成
cat <<'EOF' > /home/Competition2025/P10/P10U001/work/axolotl/axolotl_train_qwen3.sh
#!/bin/bash

#SBATCH --job-name=axolotl_train
#SBATCH --partition=P10
#SBATCH --nodes=3
#SBATCH --gpus-per-node=8
#SBATCH --gpus-per-task=8
#SBATCH --ntasks-per-node=1
#SBATCH --output=/home/Competition2025/P10/P10U001/slurm/%x_%j.log
#SBATCH --error=/home/Competition2025/P10/P10U001/slurm/%x_%j.log
#SBATCH --wait-all-nodes=1
#SBATCH --mem=0
#SBATCH --cpus-per-task=128

set -euxo pipefail

export TORCH_NCCL_ASYNC_ERROR_HANDLING=1
export NCCL_DEBUG=INFO
export NCCL_P2P_LEVEL=NVL

export NCCL_IB_DISABLE=0
export NCCL_SOCKET_IFNAME="enp25s0np0"
export NCCL_BUFFSIZE=2097152

# Hugging Face関連の設定
export HF_HOME=$HF_HOME
export HF_TOKEN=$HF_TOKEN

# export NCCL_BLOCKING_WAIT=1
# export TORCH_NCCL_BLOCKING_WAIT=1
export AXOLOTL_NCCL_TIMEOUT=7200

# W&B設定
export WANDB_API_KEY=$WANDB_API_KEY
export WANDB_ENTITY="aratako-lm"
#export WANDB_PROJECT="axolotl_training"

# 各タスクがマスターノードを見つけられるように環境変数を設定
export MASTER_ADDR=$(scontrol show hostnames $SLURM_NODELIST | head -n 1)
export MASTER_PORT=$(( 29000 + SLURM_JOB_ID % 1000 ))

# 設定ファイルのパス（ローカルまたはURL）
export CONFIG_PATH=/home/Competition2025/P10/P10U001/work/axolotl/axolotl_qwen3_235b_fsdp.yaml

# P10環境のconda
set +u 
source ~/.bashrc || true # もしくは: source /opt/conda/etc/profile.d/conda.sh 
conda activate axo-fsdp 
set -u

# --- add: custom trainer を全ノードに見せる ---
: "${PYTHONPATH:=}"
export PYTHONPATH="/home/Competition2025/P10/P10U001/work/axolotl:${PYTHONPATH:-}"

# --- add: KV/HS 蒸留の環境変数（必要なら教師モデルも）---
export AXO_KVSD_ALPHA=0.5
export AXO_KVSD_LAYERS=all
export AXO_KVSD_TOKEN_STRIDE=4
export AXO_KVSD_SHARE_STUDENT=true
export AXO_KVSD_HS_FALLBACK=true
# 任意: 別教師を使うなら
# export AXO_KVSD_TEACHER_MODEL="Qwen/Qwen2.5-7B-Instruct"
# export AXO_KVSD_TEACHER_4BIT=true

# srun --nodes=${SLURM_JOB_NUM_NODES} --ntasks-per-node=1 \
# bash -c '
# torchrun \
# --nnodes '"${SLURM_JOB_NUM_NODES}"' \
# --nproc_per_node '"${SLURM_GPUS_PER_NODE:-8}"' \
# --node_rank $SLURM_NODEID \
# --rdzv_backend c10d \
# --rdzv_id '"${SLURM_JOB_ID}"' \
# --rdzv_endpoint '"${MASTER_ADDR}:${MASTER_PORT}"' \
# -m axolotl.cli.train '"${CONFIG_PATH}"' \
# --deepspeed /home/Competition2025/P10/P10U001/work/axolotl/deepspeed_configs/zero3_bf16.json
# '

echo "Host: $(hostname)"
echo "Current user limits:"
ulimit -a
echo "Shared memory size:"
df -h /dev/shm

ulimit -v unlimited
ulimit -m unlimited

ulimit -a

srun -l --export=ALL \
bash -c '
set -eux
torchrun \
--nnodes=${SLURM_JOB_NUM_NODES} \
--nproc_per_node=${SLURM_GPUS_PER_NODE:-8} \
--node_rank=${SLURM_NODEID} \
--rdzv_backend=c10d \
--rdzv_id=${SLURM_JOB_ID} \
--rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} \
-m axolotl.cli.train ${CONFIG_PATH}
'
EOF

3.4. 実行

# 1) 必要な環境変数をこのシェルに定義（例）
export HF_HOME="$HOME/.cache/huggingface"         # 好きな場所でOK
export HF_TOKEN="hf_xxxxxxxx"                     # あなたのHFトークン
export WANDB_API_KEY="xxxxxxxxxxxxxxxxxxxxxxxx"   # あなたのW&B APIキー

JOBID=$(sbatch --parsable --export=ALL,HF_HOME,HF_TOKEN,WANDB_API_KEY,PYTHONPATH \
  /home/Competition2025/P10/P10U001/work/axolotl/axolotl_train_qwen3.sh)
echo "JOBID=$JOBID"

# ログを待ってから追う
LOG="/home/Competition2025/P10/P10U001/slurm/axolotl_train_${JOBID}.log"
until [ -f "$LOG" ] || [ -f "slurm-${JOBID}.out" ]; do sleep 1; done
tail -n +1 -F "${LOG:-slurm-${JOBID}.out}"

まとめ

Axolotl の trainer_cls × YAML × 環境変数という三点セットにより、HARI（KV/HS 自己蒸留）を巨大モデルでも安全に試行錯誤できる形に落とし込めました。
まずは AXO_KVSD_ALPHA=0.9 から始め、layers・stride・データ量を微調整して、自分のクラスターと課題に合う “ちょうどよいHARI” を見つけてください。

本プロジェクトは、国立研究開発法人新エネルギー・産業技術総合開発機構(以下「NEDO」)の「日本語版医療特化型LLMの社会実装に向けた安全性検証・実証」における基盤モデルの開発プロジェクトの一環として行われます。