松尾研 LLM開発コンペ2025 予選1位モデルの学習方法について

はじめに

私達のチームは 松尾研 LLM開発コンペ2025 の予選で1位となりました。本記事では、その際に用いたモデルの学習方法をまとめます。

予選での戦略

予選では、モデルの性能を上げることよりも、もともと HLE のスコアが高いモデルを忘却を防いで学習し、そのまま提出することを重視しました。

この戦略をとった理由

予選で使えるモデルの中で最もスコアの高い DeepSeek-R1-0528 を学習・推論するのは非常に難しいため、それだけで勝てる見込みがあったからです。

決勝で分かったこと

決勝を通じて明らかになったのは、上記の戦略に基づいて開発した、なるべく忘却せずに学習する手法が、HLE に対するスコアを再現性をもって向上させられるという点です。

HARIチューニングについて

当初このチューニングを、よく分からないため 「謎チューニング」 と呼んでいましたが、その性能を認め、現在は 「HARIチューニング」 と呼んでいます。
HARIチューニングの詳細は別のレポートでご紹介する予定です。

本記事で紹介する内容

今回は、予選モデルの学習に用いたスクリプトをご紹介します。

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_clean python=3.10

# Activate the environment
conda activate deepseek_qlora_clean

# 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. ノード確保

srun --partition P10 --nodes=1 --gpus-per-node=8 --mem=0 --cpus-per-task=240 --time=08:00:00 --pty bash -i

2.2. condaアクティベート

conda activate deepseek_qlora_clean

2.3. 環境変数設定＋ディレクトリ作成

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"

2.4. config

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

cat > qlora.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=1
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

# 方式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.5. スクリプト

cat > train_qlora.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

# ========== 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_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)

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

# ===== 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"
offload=os.path.join(nvme_dir,"offload"); 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固定＋凍結
for name, module in model.named_modules():
    if name.endswith(".mlp.gate"):
        module.eval()
        for p in module.parameters(): p.requires_grad=False

# 最後層 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:
            ds = load_dataset(repo_id, split=split_name, streaming=True,
                              token=HF_TOKEN, trust_remote_code=True)
            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()
for name, module in model.named_modules():
    if name.endswith(".mlp.gate"):
        module.eval()
        for p in module.parameters(): p.requires_grad = False

# 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

# Training bookkeeping
num_samples = len(enc["input_ids"])
num_batches_per_epoch = math.ceil(num_samples / BS)
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]

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

    for i in range(0, num_samples, BS):
        batch_inputs = {
            "input_ids": enc["input_ids"][i:i+BS],
            "attention_mask": enc["attention_mask"][i:i+BS],
        }
        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:
            with torch.no_grad():
                try:
                    ctx = model.disable_adapter()
                except AttributeError:
                    ctx = torch.no_grad()
                with ctx:
                    ref_out = model(input_ids=batch["input_ids"],
                                    attention_mask=batch["attention_mask"])
                    ref_logits = ref_out.logits

            T = KL_T if KL_T > 0 else 1.0
            s = out.logits.float() / T
            t = ref_logits.float() / 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

        loss.backward(); accum += 1

        should_step = (accum % max(GA,1) == 0)
        last_batch = (i + BS >= num_samples) and (epoch == EPOCHS-1)
        if should_step or last_batch:
            if CLIP_NORM and CLIP_NORM > 0:
                torch.nn.utils.clip_grad_norm_(trainable, CLIP_NORM)
            opt.step(); sched.step(); opt.zero_grad(set_to_none=True)
            accum = 0; step += 1
            if (step % LOG_N == 0) or last_batch:
                print(f"step {step}/{total_update_steps} loss={loss.item():.4f}", 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,
}
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.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.6. 実行

export NCCL_IB_DISABLE=1
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7

QLORA_ENV_OVERRIDE=1 QLORA_ENV_FILE=./qlora.env \
python train_qlora.py 2>&1 | tee -a train_loop.log

まとめ

本記事では、予選モデルの学習に用いたスクリプト一式（環境構築、NVMe キャッシュ設定、Slurm 実行例、qlora.env、train_qlora.py、実行コマンド）を掲載しました。
実装のキーポイントは LAMBDA_KL=0.9 による参照 KL 正則化（忘却抑制）。この設定が HARI の中核です。

以上を踏まえ、まずは記事内の手順どおり環境を再現し、qlora.env の各値（特にデータセットとトークン）を自身の環境に合わせて調整して実行してください。

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