マルチノードでDeepSeek-R1-0528 にHARIする方法について

はじめに

DeepSeek-R1-0528 をマルチノードで HARI（忘却抑制のための参照正則化つき軽量 SFT）するための実運用レシピをまとめます。

目的は、巨大 MoE を「各ノードの全 GPU をフル活用」しつつ、元モデルの振る舞いを保ったまま僅かなデータでスコアを底上げすること。SLURM 環境で 1 ノード＝1 プロセスで動かし、ノード内は device_map によるレイヤー分散、ノード間は DDP（torchrun） で同期します。
HARI の核は、LoRA を有効にした学習損失に対して、LoRA を一時無効化した“ベース参照出力”との KL 正則化（温度 (T) 付き）を加えることです。オフロード領域、データ分割、MoE のゲート固定、KL の計算負荷など、分散特有の落とし穴を潰し込んであります。

1. 環境構築

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

# Create a new conda environment with Python (you can specify the version you need)
conda create -n deepseek_qlora_mulch python=3.10

# Activate the environment
conda activate deepseek_qlora_mulch

# PyTorch/cu121
pip install --no-cache-dir --force-reinstall \
  torch==2.5.1+cu121 torchvision==0.20.1+cu121 torchaudio==2.5.1+cu121 \
  --index-url https://download.pytorch.org/whl/cu121

# bnb/triton/accelerate/transformers/peft を固定（依存を引かない）
pip install --no-cache-dir --force-reinstall --no-deps \
  triton==3.1.0 bitsandbytes==0.43.1 accelerate==0.24.1 transformers==4.36.0 peft==0.6.0

pip install huggingface_hub

# 不足している依存パッケージをインストール
pip install --no-cache-dir \
  regex \
  psutil \
  safetensors \
  tokenizers==0.15.2 \
  datasets \
  scipy

# または、特定バージョンで固定したい場合
pip install --no-cache-dir \
  regex!=2019.12.17 \
  psutil \
  safetensors>=0.3.1 \
  tokenizers==0.15.2
  
 # Pythonとcudaバージョンに合うwheelを直接インストール
pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.6.3/flash_attn-2.6.3+cu123torch2.3cxx11abiFALSE-cp310-cp310-linux_x86_64.whl

# Set environment variables
export LD_LIBRARY_PATH="$CONDA_PREFIX/lib:$LD_LIBRARY_PATH"
export BNB_CUDA_VERSION=121

2. HARI

2.1. sbatch

cat > run_qlora_mulch.sbatch <<'SB'
#!/bin/bash
#SBATCH -J qlora_mulch
#SBATCH -N 2
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=8
#SBATCH --gres=gpu:8
#SBATCH --exclusive
#SBATCH -t 48:00:00
#SBATCH -p P10
#SBATCH --export=ALL

# set -u は外す（/etc/bashrc の未定義変数対策）
set -eo pipefail

# ===== NVMe / cache =====
export NVME_DIR=/nvme12/$USER
export TMPDIR=$NVME_DIR/tmp
export HF_HOME=$NVME_DIR/hf_home
export TRANSFORMERS_CACHE=$NVME_DIR/hf_cache
export CUDA_CACHE_PATH=$NVME_DIR/ComputeCache
mkdir -p "$TMPDIR" "$HF_HOME" "$TRANSFORMERS_CACHE" "$CUDA_CACHE_PATH"

# ===== runtime =====
export BASHRCSOURCED=1
source ~/.bashrc
conda activate deepseek_qlora_mulch

export QLORA_ENV_OVERRIDE=1
export QLORA_ENV_FILE=./qlora_mulch.env
export BNB_CUDA_VERSION=121
export NCCL_DEBUG=WARN
export NCCL_ASYNC_ERROR_HANDLING=1
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export OMP_NUM_THREADS=8
# 追加で詳細ログが欲しいときはコメント解除
# export TORCH_DISTRIBUTED_DEBUG=DETAIL
# export TORCH_SHOW_CPP_STACKTRACES=1
ulimit -n 131072

# IB/ネットワークIF自動判定（ib0が無ければ lo,docker0 を除外）
export NCCL_SOCKET_IFNAME=ib0
if ! ip -br a | awk '{print $1}' | grep -qx "$NCCL_SOCKET_IFNAME"; then
  echo "[WARN] $NCCL_SOCKET_IFNAME not found; falling back to ^lo,docker0"
  export NCCL_SOCKET_IFNAME='^lo,docker0'
fi

# rendezvous
HOSTS=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
MASTER_NODE=$(echo "$HOSTS" | head -n1)
MASTER_ADDR=$(getent ahostsv4 "$MASTER_NODE" | awk 'NR==1{print $1}')
MASTER_PORT=$((15000 + (${SLURM_JOB_ID} % 10000) ))

echo "MASTER_ADDR=${MASTER_ADDR} MASTER_PORT=${MASTER_PORT} NODES=${SLURM_NNODES}"

# ==== ここがポイント：1ノード=1プロセス ====
srun --nodes=${SLURM_NNODES} --ntasks-per-node=1 --export=ALL --kill-on-bad-exit=1 bash -lc '
  export BASHRCSOURCED=1
  source ~/.bashrc
  conda activate deepseek_qlora_mulch
  export MASTER_ADDR='"$MASTER_ADDR"'
  export MASTER_PORT='"$MASTER_PORT"'
  export NODE_RANK=${SLURM_NODEID}
  echo "[NODE ${SLURM_NODEID}] torchrun with nproc_per_node=1 (model uses ALL GPUs on the node)"
  torchrun --nproc_per_node=1 \
           --nnodes='"$SLURM_NNODES"' \
           --node_rank=${NODE_RANK} \
           --rdzv_backend=c10d \
           --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} \
           train_qlora_mulch.py
'
SB

2.2. config

LAMBDA_KL=0.9にして忘却防ぐのがポイントです。これがHARIです。

cat > qlora_mulch.env <<'ENV'
# ===== QLoRA runtime config (.env) =====
MODEL_LOCAL_DIR=/nvme12/P10U001/deepseek_src
NVME_DIR=/nvme12/$USER

# HF private datasets 用
HF_TOKEN=

# 乱数シード
SEED=42

# 学習ノブ
BS=2
EPOCHS=10
GRAD_ACCUM=2
MAX_LEN=4096
CLIP_NORM=1.0
WARMUP_RATIO=0.03
LOG_INTERVAL=100
GRAD_CP=1

# 参照KL正則化（0で無効）
LAMBDA_KL=0.9
KL_T=1.5

# ===== データセット設定（新規） =====
# 方式A: 文字列で簡易指定
#   repo[:limit][@split+split], 複数はカンマ区切り。limit=ALL/未指定で全件。
SFT_DATASETS=oNo-1/MedMCQA:ALL@train,oNo-1/OlympiadBench:8000@train

# 方式B: JSONで厳密指定（Aよりこちらが優先）
# SFT_DATASETS_JSON='[{"name":"oNo-1/MedMCQA","max_n":null,"splits":["train"]},{"name":"oNo-1/OlympiadBench","max_n":8000,"splits":["train"]}]'

# 方式C: 総件数の上限（全体カット）。未指定なら無制限
# SFT_LIMIT=12000

# デフォルトで試す split（個別指定がなければこれを使用）
SFT_SPLITS=train,validation,test

# ===== Config loader の振る舞い =====
QLORA_ENV_OVERRIDE=1
ENV

2.3. スクリプト

cat > train_qlora_mulch.py <<'PY'
import os, re, torch, math, json
from datetime import datetime
try:
    from zoneinfo import ZoneInfo
    JST = ZoneInfo("Asia/Tokyo")
except Exception:
    JST = None

from transformers import (
    AutoConfig, AutoModelForCausalLM, AutoTokenizer,
    BitsAndBytesConfig, DataCollatorForLanguageModeling,
    get_linear_schedule_with_warmup
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from bitsandbytes.optim import PagedAdamW8bit
torch.backends.cuda.matmul.allow_tf32 = True
import torch.nn.functional as F  # KL
import torch.distributed as dist  # ★ 追加
import socket  # ★ 追加: ノード名取得

# ========== 分散初期化（最小追加） ==========
def _dist_init_if_needed():
    if dist.is_available() and not dist.is_initialized() and int(os.environ.get("WORLD_SIZE","1")) > 1:
        dist.init_process_group(backend="nccl")
    world = dist.get_world_size() if dist.is_initialized() else 1
    rank = dist.get_rank() if dist.is_initialized() else 0
    local_rank = int(os.environ.get("LOCAL_RANK","0"))
    if torch.cuda.is_available():
        torch.cuda.set_device(local_rank)
    return world, rank, local_rank

# ========== Config loader (.env / JSON) ==========
def _load_env_from_file(path, override=False):
    def _strip_inline_comment(s: str) -> str:
        in_single = False; in_double = False; out = []
        for ch in s:
            if ch == "'" and not in_double: in_single = not in_single
            elif ch == '"' and not in_single: in_double = not in_double
            elif ch == '#' and not in_single and not in_double: break
            out.append(ch)
        return ''.join(out)

    try:
        with open(path, "r", encoding="utf-8") as f:
            for raw in f:
                line = raw.strip()
                if not line or line.startswith("#") or "=" not in line: continue
                k, v = line.split("=", 1)
                k = k.strip()
                v = _strip_inline_comment(v).strip()
                if (v.startswith('"') and v.endswith('"')) or (v.startswith("'") and v.endswith("'")):
                    v = v[1:-1]
                v = os.path.expandvars(v)
                if not override and k in os.environ: continue
                os.environ[k] = v
        print(f"🔧 loaded .env: {path} (override={override})", flush=True)
    except Exception as e:
        print(f"⚠️ QLORA_ENV_FILE load failed: {type(e).__name__}: {e}", flush=True)

def _load_env_from_json_text(text, override=False):
    try:
        data = json.loads(text)
        if not isinstance(data, dict): raise ValueError("top-level is not an object")
        for k, v in data.items():
            if v is None: continue
            val = os.path.expandvars(str(v))
            if not override and k in os.environ: continue
            os.environ[k] = val
        print(f"🔧 loaded JSON config (override={override})", flush=True)
    except Exception as e:
        print(f"⚠️ QLORA_CONFIG JSON load failed: {type(e).__name__}: {e}", flush=True)

_override = os.getenv("QLORA_ENV_OVERRIDE", "0") == "1"
# ★ 追加: デフォルトの .env を自動ロード
if not os.getenv("QLORA_ENV_FILE") and os.path.isfile("./qlora_mulch.env"):
    os.environ["QLORA_ENV_FILE"] = "./qlora_mulch.env"
_env_file = os.getenv("QLORA_ENV_FILE")
if _env_file and os.path.isfile(_env_file):
    _load_env_from_file(_env_file, _override)
_cfg_json_file = os.getenv("QLORA_CONFIG_JSON")
if _cfg_json_file and os.path.isfile(_cfg_json_file):
    with open(_cfg_json_file, "r", encoding="utf-8") as _f:
        _load_env_from_json_text(_f.read(), _override)
elif os.getenv("QLORA_CONFIG"):
    _load_env_from_json_text(os.getenv("QLORA_CONFIG"), _override)
# =================================================

# ===== Env knobs =====
SEED         = int(os.environ.get("SEED","42"))
BS           = int(os.environ.get("BS","1"))
EPOCHS       = int(os.environ.get("EPOCHS","1"))
GA           = int(os.environ.get("GRAD_ACCUM","1"))
CLIP_NORM    = float(os.environ.get("CLIP_NORM","1.0"))
WARMUP_RATIO = float(os.environ.get("WARMUP_RATIO","0.03"))
LOG_N        = max(1, int(os.environ.get("LOG_INTERVAL","50")))
MAX_LEN      = int(os.environ.get("MAX_LEN","8192"))
HF_TOKEN     = os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")
GRAD_CP      = os.getenv("GRAD_CP","0") == "1"
# 参照KL
LAMBDA_KL    = float(os.environ.get("LAMBDA_KL", "0.0"))
KL_T         = float(os.environ.get("KL_T", "1.0"))
# データセット設定（新規）
SFT_LIMIT_ENV = os.environ.get("SFT_LIMIT", None)  # 既存の総上限（任意）
SFT_DATASETS  = os.environ.get("SFT_DATASETS")     # 例: "oNo-1/MedMCQA:ALL@train,oNo-1/OlympiadBench:8000"
SFT_DATASETS_JSON = os.environ.get("SFT_DATASETS_JSON")  # JSON: [{"name":"...", "max_n":null, "splits":["train"]}, ...]
SFT_SPLITS_DEF = os.environ.get("SFT_SPLITS", "train,validation,test")  # デフォルトで試すsplit

# full seeding
import random, numpy as np
random.seed(SEED); np.random.seed(SEED)
torch.manual_seed(SEED); torch.cuda.manual_seed_all(SEED)

# ★ 追加：分散情報
WORLD_SIZE, RANK, LOCAL_RANK = _dist_init_if_needed()
HOSTNAME = socket.gethostname()  # ★ 追加: このランクのノード名

def _mlp_forward_identity(self, hidden_states, *args, **kwargs):
    return hidden_states

# ★ 追加: gate を常に eval 固定＆凍結するユーティリティ
def _set_gates_eval_and_freeze(m):
    core = m.module if hasattr(m, "module") else m
    for name, module in core.named_modules():
        if name.endswith(".gate") or ("gate" in module.__class__.__name__.lower()):
            module.eval()
            for p in module.parameters():
                p.requires_grad = False

# ===== Dataset config parsing (new) =====
def _parse_default_splits():
    return [s.strip() for s in SFT_SPLITS_DEF.split(",") if s.strip()]

def _parse_datasets_plan():
    # 1) JSON優先
    if SFT_DATASETS_JSON:
        plan = []
        try:
            arr = json.loads(SFT_DATASETS_JSON)
            for x in arr:
                if isinstance(x, dict):
                    name = x.get("name") or x.get("repo") or x.get("id")
                    if not name: continue
                    max_n = x.get("max_n", None)
                    splits = x.get("splits", None)
                    if splits is not None and isinstance(splits, str):
                        splits = [s.strip() for s in splits.split(",") if s.strip()]
                    plan.append((name, max_n, splits))
                elif isinstance(x, (list, tuple)) and len(x)>=1:
                    name = x[0]; max_n = x[1] if len(x)>1 else None
                    splits = x[2] if len(x)>2 else None
                    plan.append((name, max_n, splits))
        except Exception as e:
            print(f"⚠️ SFT_DATASETS_JSON parse failed: {type(e).__name__}: {e}", flush=True)
        if plan: return plan
    # 2) 文字列 "repo[:limit][@split+split],repo2[:limit],..."
    if SFT_DATASETS:
        plan = []
        for part in [p.strip() for p in SFT_DATASETS.split(",") if p.strip()]:
            name = part; max_n=None; splits=None
            if "@" in part:
                base, sp = part.split("@",1)
                name = base
                splits = [s.strip() for s in sp.replace("+",",").split(",") if s.strip()]
            if ":" in name:
                nm, lim = name.split(":",1)
                name = nm.strip()
                lim = lim.strip().lower()
                if lim and lim not in ("all","none"):
                    try: max_n=int(lim)
                    except: max_n=None
                else:
                    max_n=None
            plan.append((name, max_n, splits))
        if plan: return plan
    # 3) デフォルト（従来通り）
    return [("oNo-1/MedMCQA", None, None), ("oNo-1/OlympiadBench", 8000, None)]

model_path=os.environ["MODEL_LOCAL_DIR"]; nvme_dir=os.environ["NVME_DIR"]
assert model_path and nvme_dir

cfg=AutoConfig.from_pretrained(model_path, trust_remote_code=True)
num_layers=cfg.num_hidden_layers
gpus=list(range(torch.cuda.device_count()))
device_map={"model.embed_tokens":"cuda:0"}
for i in range(num_layers): device_map[f"model.layers.{i}"]=f"cuda:{i%len(gpus)}"
device_map["model.norm"]=f"cuda:{gpus[-1]}"; device_map["lm_head"]=f"cuda:{gpus[-1]}"

# QLoRA: 4-bit NF4 + k-bit準備
bnb=BitsAndBytesConfig(load_in_4bit=True,bnb_4bit_use_double_quant=True,
                       bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16)

max_memory={f"cuda:{i}":"78GiB" for i in gpus}; max_memory["cpu"]="300GiB"
# ★ 変更(A): オフロード先をランク別に分離（envのRANKではなく実ランクを使用）
offload=os.path.join(nvme_dir, f"offload_rank{RANK}"); os.makedirs(offload,exist_ok=True)

tok=AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, use_fast=True)
if tok.pad_token is None and tok.eos_token is not None: tok.pad_token=tok.eos_token
tok.padding_side="right"; tok.truncation_side="right"

model=AutoModelForCausalLM.from_pretrained(
    model_path, trust_remote_code=True, quantization_config=bnb,
    device_map=device_map, low_cpu_mem_usage=True, offload_state_dict=True,
    offload_folder=offload, max_memory=max_memory, torch_dtype=torch.bfloat16,
)
model.config.use_cache=False
model=prepare_model_for_kbit_training(model)
if GRAD_CP:
    model.gradient_checkpointing_enable()

# MoE gate: eval固定＋凍結（ユーティリティ経由）
_set_gates_eval_and_freeze(model)

# 最後層 self_attn の線形候補を自動検出（Q優先）
last_idx=num_layers-1; cand=set()
for n,m in model.named_modules():
    if f"model.layers.{last_idx}.self_attn." in n and "linear" in m.__class__.__name__.lower():
        cand.add(n.split(".")[-1])
prefs=[["q_a_proj","q_b_proj"],["kv_b_proj"],["q_proj","v_proj"],["wq","wv"],["qkv_proj"],["query_key_value"],["o_proj","wo"]]
targets=[]
for g in prefs:
    hit=[x for x in g if x in cand]
    if hit: targets=hit; break
if not targets: targets=sorted(cand)[:1] or ["o_proj"]
print("candidate attn linear:", sorted(cand))
print("use target_modules:", targets)

# LoRA（QLoRAのAdapter部）。最後層以外のLoRAは凍結
lora_cfg=LoraConfig(r=4, lora_alpha=8, lora_dropout=0.0,
                    target_modules=targets, bias="none", task_type="CAUSAL_LM")
model=get_peft_model(model,lora_cfg)
for n,m in model.named_modules():
    if "lora_" in n and (f"model.layers.{last_idx}." not in n):
        for p in getattr(m,"parameters",lambda:[])(): p.requires_grad=False

# ===================== データ部分（SFTテンプレ適用） =====================
from datasets import load_dataset
import string

def build_user_text(ex):
    q = str(ex.get("question","")).strip()
    opts = ex.get("choices") or ex.get("options") or None
    if isinstance(opts, list) and len(opts) > 0:
        abc = list(string.ascii_uppercase)
        lines = [f"{abc[i]}. {str(o)}" for i,o in enumerate(opts)]
        return f"{q}\n\n選択肢:\n" + "\n".join(lines)
    return q

def build_assistant_text(ex):
    ans = str(ex.get("answer","")).strip()
    cot = ex.get("cot") or ex.get("rationale") or ex.get("explanation")
    if cot:
        cot = str(cot).strip()
        return f"<think>\n{cot}\n</think>\n<answer>\n{ans}\n</answer>"
    else:
        return f"<answer>\n{ans}\n</answer>"

def to_messages(ex):
    return [
        {"role":"system","content":"You are DeepSeek-R1."},
        {"role":"user","content": build_user_text(ex)},
        {"role":"assistant","content": build_assistant_text(ex)},
    ]

DEFAULT_SPLITS = _parse_default_splits()
RAW_PLAN = _parse_datasets_plan()  # list of (name, max_n, splits|None)

# Arrow の型キャストエラー回避：ストリーミング（HF token明示）
def _stream_examples(repo_id, splits=None):
    splits = splits or DEFAULT_SPLITS
    for split_name in splits:
        try:
            # ★ 修正: datasets の trust_remote_code を削除
            ds = load_dataset(repo_id, split=split_name, streaming=True, token=HF_TOKEN)
            for ex in ds:
                yield ex
        except Exception as e:
            print(f"⚠️ stream skip {repo_id}:{split_name} -> {type(e).__name__}: {e}", flush=True)

texts = []; name_fragments = []; resolved_datasets = []
limit = int(SFT_LIMIT_ENV) if SFT_LIMIT_ENV is not None else None

for name, max_n, splits in RAW_PLAN:
    added = 0
    for ex in _stream_examples(name, splits):
        if limit is not None and len(texts) >= limit: break
        if max_n is not None and added >= int(max_n): break
        messages = to_messages(ex)
        text = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
        texts.append(text); added += 1
    frag = f"{name}[n{added}" + (f"/{max_n}]" if max_n is not None else "]")
    name_fragments.append(frag)
    resolved_datasets.append({
        "name": name,
        "max_n": max_n,
        "splits": splits or DEFAULT_SPLITS,
        "loaded": added
    })

SFT_DATASET_NAME = " + ".join(name_fragments)
SFT_SAMPLE_COUNT = len(texts)
if SFT_SAMPLE_COUNT == 0:
    raise RuntimeError("No samples loaded. Check HF token/permissions or dataset availability.")

# アシスタントのみ損失のカット点（トークン位置）
assistant_prefix_token_lens = []
for t in texts:
    cut_char = t.find("<think>")
    if cut_char == -1: cut_char = t.find("<answer>")
    if cut_char == -1: cut_char = 0
    enc_one = tok(t, truncation=True, max_length=MAX_LEN, padding=False,
                  add_special_tokens=True, return_offsets_mapping=True)
    offsets = enc_one["offset_mapping"]; cut_tok = len(offsets)
    for j, off in enumerate(offsets):
        if off is None: continue
        start = off[0] if isinstance(off, (list, tuple)) else None
        if start is not None and start >= cut_char:
            cut_tok = j; break
    assistant_prefix_token_lens.append(cut_tok)

# 動的パディングでエンコード
enc = tok(texts, truncation=True, max_length=MAX_LEN, padding=False)
collator = DataCollatorForLanguageModeling(tokenizer=tok, mlm=False)

# 初回シャッフル（再現性は SEED で制御）
order = list(range(len(enc["input_ids"]))); random.shuffle(order)
for k in list(enc.keys()): enc[k] = [enc[k][i] for i in order]
assistant_prefix_token_lens = [assistant_prefix_token_lens[i] for i in order]
# ===================== データ部分ここまで =====================

# optimizer（LoRAパラメータのみ）＋ scheduler
trainable=[p for p in model.parameters() if p.requires_grad]
opt=PagedAdamW8bit(trainable, lr=5e-5)

model.train()

# Re-freeze MoE gate after model.train()（ユーティリティで再適用）
_set_gates_eval_and_freeze(model)

# MLPはforward恒等にしない（品質維持）。パラメータは凍結のみ
for name, module in model.named_modules():
    if name.endswith(".mlp"):
        for p in module.parameters(): p.requires_grad = False

# DeepSeek MoE shared_experts の一時回避: 例外時のみ恒等フォールバック
from types import MethodType
def _wrap_mlp_forward_with_fallback(module):
    orig_forward = module.forward
    def _safe_forward(self, hidden_states, *args, **kwargs):
        try:
            return orig_forward(hidden_states, *args, **kwargs)
        except UnboundLocalError as e:
            if "y" in str(e):  # 'y referenced before assignment'
                return hidden_states
            raise
    module.forward = MethodType(_safe_forward, module)
for name, m in model.named_modules():
    if name.endswith(".mlp"): _wrap_mlp_forward_with_fallback(m)

first_device=model.get_input_embeddings().weight.device

# ★ 追加：DDP ラップ（ノード内複数GPUに跨るので device_ids=None）
if dist.is_available() and dist.is_initialized():
    from torch.nn.parallel import DistributedDataParallel as DDP
    model = DDP(model, device_ids=None, broadcast_buffers=False, find_unused_parameters=True)
    # DDP後も gate を eval に維持
    _set_gates_eval_and_freeze(model)

# Training bookkeeping（各ランクのステップ数で再計算）
num_samples = len(enc["input_ids"])
num_batches_per_epoch = math.ceil(num_samples / (BS * max(int(os.environ.get("WORLD_SIZE","1")),1)))
total_update_steps = math.ceil((num_batches_per_epoch * EPOCHS) / max(GA,1))
warmup_steps = int(WARMUP_RATIO * total_update_steps)
sched = get_linear_schedule_with_warmup(opt, warmup_steps, total_update_steps)

step = 0  # optimizer steps
accum = 0
opt.zero_grad(set_to_none=True)

def _epoch_reshuffle():
    global enc, assistant_prefix_token_lens
    order2 = list(range(len(enc["input_ids"]))); random.shuffle(order2)
    for k in list(enc.keys()): enc[k] = [enc[k][i] for i in order2]
    assistant_prefix_token_lens = [assistant_prefix_token_lens[i] for i in order2]

# ★ 追加(B): drop_last相当の上限
global_pairs = (num_samples // (BS * max(WORLD_SIZE,1))) * (BS * max(WORLD_SIZE,1))

# Train loop（エポック駆動／勾配累積／動的パディング）
for epoch in range(EPOCHS):
    if epoch > 0: _epoch_reshuffle()

    # ★ ランクごとにバッチをストライド分割（drop_last相当で global_pairs まで）
    for i in range(RANK*BS, global_pairs, BS*max(WORLD_SIZE,1)):
        batch_inputs = {
            "input_ids": enc["input_ids"][i:i+BS],
            "attention_mask": enc["attention_mask"][i:i+BS],
        }
        if len(batch_inputs["input_ids"]) == 0:
            continue

        batch = collator([
            {"input_ids": x, "attention_mask": y}
            for x, y in zip(batch_inputs["input_ids"], batch_inputs["attention_mask"])
        ])
        for k in list(batch.keys()):
            batch[k] = batch[k].to(first_device, non_blocking=True)

        # アシスタントのみ損失（<think>/<answer> 以前をマスク）
        labels = batch["labels"]; seq_len = labels.size(1)
        for b_idx, cut in enumerate(assistant_prefix_token_lens[i:i+BS]):
            c = min(cut, seq_len)
            if c > 0: labels[b_idx, :c] = -100
        batch["labels"] = labels

        out=model(**batch); loss=out.loss

        # 参照KL（LoRA無効=ベース）との距離を加算（LAMBDA_KL>0のときのみ）
        if LAMBDA_KL > 0.0:
            # teacher側は一時的にeval()にしてから計算
            model_core = (model.module if hasattr(model,"module") else model)
            was_train = model_core.training
            try:
                model_core.eval()
                with torch.no_grad():
                    try:
                        ctx = model_core.disable_adapter()
                    except AttributeError:
                        ctx = torch.no_grad()
                    with ctx:
                        ref_out = model_core(
                            input_ids=batch["input_ids"],
                            attention_mask=batch["attention_mask"]
                        )
                        ref_logits = ref_out.logits
            finally:
                if was_train:
                    model_core.train()
                    # ★ 再固定: gate を再び eval
                    _set_gates_eval_and_freeze(model_core)

            # ★ 変更: KL計算を bf16/half で行い、メモリを半減
            T = KL_T if KL_T > 0 else 1.0
            dtype_kl = torch.bfloat16 if out.logits.dtype == torch.bfloat16 else torch.float16
            s = (out.logits.to(dtype_kl) / T)
            t = (ref_logits.to(dtype_kl) / T)
            logps = F.log_softmax(s, dim=-1)
            probs_t = F.softmax(t, dim=-1)
            kl_per_tok = F.kl_div(logps, probs_t, reduction='none').sum(-1)  # (B, S)
            mask = batch["attention_mask"]
            kl_mean = (kl_per_tok * mask).sum() / mask.sum().clamp(min=1)
            loss = loss + (LAMBDA_KL * (T * T)) * kl_mean

            # ★ 後片付けでピークメモリを下げる
            del ref_out, ref_logits, s, t, logps, probs_t, kl_per_tok
            torch.cuda.empty_cache()

        loss.backward(); accum += 1

        should_step = (accum % max(GA,1) == 0)
        # 端数落としに合わせて最終バッチ判定もglobal_pairs基準に
        last_batch = (i + BS*max(WORLD_SIZE,1) >= global_pairs) and (epoch == EPOCHS-1)
        if should_step or last_batch:
            if CLIP_NORM and CLIP_NORM > 0:
                torch.nn.utils.clip_grad_norm_([p for p in (model.module if hasattr(model,"module") else model).parameters() if p.requires_grad], CLIP_NORM)
            opt.step(); sched.step(); opt.zero_grad(set_to_none=True)
            accum = 0; step += 1
            # 学習中のログはrank0のみ
            if (((step % LOG_N) == 0) or last_batch) and (not dist.is_initialized() or dist.get_rank() == 0):
                print(f"step {step}/{total_update_steps} loss={loss.item():.4f}", flush=True)

# 出力（rank 同期→ノード一覧収集→rank0のみ保存）
if dist.is_available() and dist.is_initialized():
    dist.barrier()

# 全ノード名を収集してrank0で表示
nodes_collected = None
if dist.is_available() and dist.is_initialized():
    try:
        gathered = [None] * dist.get_world_size()
        dist.all_gather_object(gathered, HOSTNAME)
        if dist.get_rank() == 0:
            nodes_collected = sorted(set(gathered))
    except Exception as e:
        if dist.get_rank() == 0:
            print(f"⚠️ node gather failed: {type(e).__name__}: {e}", flush=True)

now = datetime.now(JST) if JST else datetime.now()
ts_disp = now.strftime("%Y-%m-%d %H:%M:%S %Z") if JST else now.strftime("%Y-%m-%d %H:%M:%S")
ts_compact = now.strftime("%Y%m%d-%H%M%S")

dataset_slug = re.sub(r'[^A-Za-z0-9._-]+','-', SFT_DATASET_NAME.replace('/','-')).strip('-_') or "dataset"
outdir_name = f"qlora_{dataset_slug}__n{SFT_SAMPLE_COUNT}__{ts_compact}"
outdir = os.path.join(nvme_dir, outdir_name); os.makedirs(outdir, exist_ok=True)

# ====== config snapshot を保存（HF_TOKENはマスク） ======
def _snapshot_env():
    keys = [
        # core
        "MODEL_LOCAL_DIR","NVME_DIR","SEED","BS","EPOCHS","GRAD_ACCUM","MAX_LEN",
        "CLIP_NORM","WARMUP_RATIO","LOG_INTERVAL","GRAD_CP",
        # KD
        "LAMBDA_KL","KL_T",
        # dataset
        "SFT_LIMIT","SFT_DATASETS","SFT_DATASETS_JSON","SFT_SPLITS",
        # loader behavior
        "QLORA_ENV_FILE","QLORA_CONFIG_JSON","QLORA_ENV_OVERRIDE",
    ]
    snap = {}
    for k in keys:
        if k in os.environ: snap[k] = os.environ[k]
    # mask secrets if present
    if "HF_TOKEN" in os.environ: snap["HF_TOKEN"] = "***"
    if "HUGGING_FACE_HUB_TOKEN" in os.environ: snap["HUGGING_FACE_HUB_TOKEN"] = "***"
    return snap

config_snapshot = {
    "env": _snapshot_env(),
    "resolved_datasets": resolved_datasets,
    "dataset_name_fragments": name_fragments,
    "samples": SFT_SAMPLE_COUNT,
    "timestamp": ts_disp,
}

if (not dist.is_available()) or (not dist.is_initialized()) or (dist.get_rank()==0):  # rank0 だけ保存
    # 保存ノード（rank0のホスト名）と参加ノード一覧をプリント
    print(f"💾 saving on host={HOSTNAME} (rank=0)", flush=True)
    if nodes_collected is not None:
        print(f"🖥️ participating nodes: {nodes_collected}", flush=True)

    with open(os.path.join(outdir, "config_snapshot.json"), "w", encoding="utf-8") as f:
        json.dump(config_snapshot, f, ensure_ascii=False, indent=2)

    # モデルと tokenizer の保存
    (model.module if hasattr(model,"module") else model).save_pretrained(outdir); tok.save_pretrained(outdir)

    # ラン情報の保存
    info = {
      "timestamp": ts_disp,
      "dataset": SFT_DATASET_NAME,
      "samples": SFT_SAMPLE_COUNT,
      "epochs": EPOCHS,
      "bs": BS,
      "grad_accum": GA,
      "max_len": MAX_LEN,
    }
    with open(os.path.join(outdir, "sft_run_info.json"), "w", encoding="utf-8") as f:
        json.dump(info, f, ensure_ascii=False, indent=2)

    print("✅ QLoRA loop done ->", outdir)
    print(f"🕒 {ts_disp} | 📚 {SFT_DATASET_NAME} | 🔢 samples={SFT_SAMPLE_COUNT} | "
          f"🧮 steps={step}/{total_update_steps}")
PY

2.4. 実行

export BASHRCSOURCED=1
sbatch run_qlora_mulch.sbatch

まとめ

本稿のレシピは、巨大 MoE を壊さず・速く・再現良く底上げするために、マルチノードを「1 ノード＝1 プロセス」、ノード内は device_map による層分散、ノード間は DDP で同期する――という最小設計に絞りました。学習自体は LoRA 有効の SFT 損失に対し、LoRA を一時無効化したベース参照出力との温度付き KLを加える HARI で忘却を抑制。あわせて MoE ゲートの eval 固定と凍結、rank 別オフロード、アシスタント部分のみの損失化、KL を bf16/half で計算してメモリを抑えるなど、分散特有の落とし穴を潰し込んでいます。

運用の第一歩としては、公開の sbatch/.env/train.py をそのまま流し、データセット差し替えと LAMBDA_KL の調整だけで十分です（安定重視なら 0.05–0.2 から開始し、必要に応じて上げる）。サンプル数は 500〜1 万程度・1 エポックが立ち上がりが早く、成果が見えたらバッチ・長さ・エポックを微調整してください。

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