recipe/sft.py

"""
AHA-Qwen3 SFT training script.

Loads pretrained Qwen3-0.6B, converts to AHA-Qwen3 (adds gate to q_proj),
and trains on am-distilled data. Only gate weights are randomly initialized;
all other weights come from the pretrained model.
"""

import json
import math
import os
import shutil
import socket
import subprocess
from typing import Optional

import datasets
import torch
import torch.distributed as dist
from trl import SFTConfig, SFTTrainer
from transformers import AutoTokenizer, TrainerCallback
from transformers.trainer_utils import get_last_checkpoint

from modeling_aha_qwen3 import (
    AHA_ROUTER_GRANULARITY,
    AHAQwen3Config,
    AHAQwen3ForCausalLM,
    aha_router_output_size,
)
from router_training_utils import RowWiseAdamW, configure_gate_only


class AHASFTTrainer(SFTTrainer):
    """Accumulate CE vs gate-aux vs distill breakdown and gate density; log with HF `loss` (total)."""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # ``AHAQwen3ForCausalLM.forward`` accepts ``**kwargs`` for attention
        # backends, but its CE loss is an ordinary local token mean and does
        # not consume Trainer's ``num_items_in_batch``.  Transformers 4.57
        # otherwise mistakes the variadic signature for a globally normalized
        # loss, multiplies it by world size, and skips the normal gradient-
        # accumulation division.  Mark the actual loss contract explicitly so
        # lr=3e-5 has the same meaning at every world size.
        self.model_accepts_loss_kwargs = False
        self._aha_ce_sum = 0.0
        self._aha_gate_aux_sum = 0.0
        self._aha_distill_sum = 0.0
        self._aha_gate_soft_sum = 0.0
        self._aha_gate_hard_sum = 0.0
        self._aha_metric_count = 0
        # Cached for loss_total aggregation in log().
        self._aha_distill_weight = float(
            getattr(self.model.config, "aha_distill_weight", 0.0)
        )
        self._aha_ce_weight = float(
            getattr(self.model.config, "aha_ce_weight", 1.0)
        )

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        loss, outputs = super().compute_loss(
            model, inputs, return_outputs=True, num_items_in_batch=num_items_in_batch
        )
        if getattr(outputs, "ce_loss", None) is not None:
            self._aha_metric_count += 1
            self._aha_ce_sum += float(outputs.ce_loss.detach().float().mean().cpu())
            if getattr(outputs, "gate_aux_loss", None) is not None:
                self._aha_gate_aux_sum += float(outputs.gate_aux_loss.detach().float().mean().cpu())
            if getattr(outputs, "distill_loss", None) is not None:
                self._aha_distill_sum += float(outputs.distill_loss.detach().float().mean().cpu())
            if getattr(outputs, "gate_soft_mean", None) is not None:
                self._aha_gate_soft_sum += float(outputs.gate_soft_mean.detach().float().mean().cpu())
            if getattr(outputs, "gate_hard_mean", None) is not None:
                self._aha_gate_hard_sum += float(outputs.gate_hard_mean.detach().float().mean().cpu())
        if return_outputs:
            return (loss, outputs)
        return loss

    def _distributed_mean_of_sums(self, sum_val: float, count: int) -> float:
        if count <= 0:
            return float("nan")
        device = self.accelerator.device
        if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1:
            t = torch.tensor([sum_val, float(count)], device=device, dtype=torch.float64)
            dist.all_reduce(t, op=dist.ReduceOp.SUM)
            return (t[0] / t[1]).item()
        return sum_val / float(count)

    def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> None:
        if self._aha_metric_count > 0:
            n = self._aha_metric_count
            ce_m = self._distributed_mean_of_sums(self._aha_ce_sum, n)
            aux_m = self._distributed_mean_of_sums(self._aha_gate_aux_sum, n)
            distill_m = self._distributed_mean_of_sums(self._aha_distill_sum, n)
            gs_m = self._distributed_mean_of_sums(self._aha_gate_soft_sum, n)
            gh_m = self._distributed_mean_of_sums(self._aha_gate_hard_sum, n)
            if not math.isnan(ce_m):
                logs["ce_loss"] = round(ce_m, 4)
            if not math.isnan(aux_m):
                logs["gate_aux_loss"] = round(aux_m, 6)
            if not math.isnan(distill_m) and distill_m != 0.0:
                logs["distill_loss"] = round(distill_m, 6)
            # loss_total = aha_ce_weight * ce + gate_aux
            #              + aha_distill_weight * distill_loss
            # (distill_m / ce_m are unweighted; apply weights here so the
            # sum matches the scalar actually added to total loss.)
            if not math.isnan(ce_m) and not math.isnan(aux_m):
                lt = self._aha_ce_weight * ce_m + aux_m
                if not math.isnan(distill_m):
                    lt = lt + self._aha_distill_weight * distill_m
                logs["loss_total"] = round(lt, 4)
                if "loss" in logs and lt > 1e-8:
                    # With ``model_accepts_loss_kwargs=False``, HF's logging
                    # window and the raw model loss should agree up to normal
                    # batch-to-batch weighting differences. A large integer
                    # ratio is a regression in Trainer loss normalization.
                    logs["hf_loss_ratio"] = round(logs["loss"] / lt, 2)
            if not math.isnan(gs_m):
                logs["gate_soft_mean"] = round(gs_m, 4)
            if not math.isnan(gh_m):
                logs["gate_hard_mean"] = round(gh_m, 4)
            self._aha_ce_sum = 0.0
            self._aha_gate_aux_sum = 0.0
            self._aha_distill_sum = 0.0
            self._aha_gate_soft_sum = 0.0
            self._aha_gate_hard_sum = 0.0
            self._aha_metric_count = 0
        return super().log(logs, start_time)


class AHAManifestCallback(TrainerCallback):
    """Copy the immutable run manifest into each Trainer checkpoint."""

    def __init__(self, manifest_path: str):
        self.manifest_path = manifest_path

    def on_save(self, args, state, control, **kwargs):
        if state.is_world_process_zero and os.path.exists(self.manifest_path):
            checkpoint_dir = os.path.join(
                args.output_dir, f"checkpoint-{state.global_step}"
            )
            if os.path.isdir(checkpoint_dir):
                shutil.copyfile(
                    self.manifest_path,
                    os.path.join(checkpoint_dir, "aha_training_manifest.json"),
                )
        return control


def get_optional_int(name: str) -> Optional[int]:
    value = os.environ.get(name)
    return None if value in {None, ""} else int(value)


def resolve_resume_checkpoint(output_dir: str) -> Optional[str]:
    resume_from_checkpoint = os.environ.get("RESUME_FROM_CHECKPOINT")
    if resume_from_checkpoint in {None, ""}:
        return None
    if resume_from_checkpoint == "latest":
        if not os.path.isdir(output_dir):
            return None
        return get_last_checkpoint(output_dir)
    return resume_from_checkpoint


def parse_fsdp_options(name: str) -> list[str]:
    value = os.environ.get(name, "").replace(",", " ").strip()
    return [item for item in value.split() if item]


def parse_bool(name: str, default: bool = False) -> bool:
    value = os.environ.get(name)
    if value is None:
        return default
    return value.lower() in {"1", "true", "yes", "on"}


def parse_report_to() -> list[str]:
    """Comma-separated integrations, e.g. REPORT_TO=wandb or wandb,tensorboard. Empty / none / off -> []."""
    raw = os.environ.get("REPORT_TO", "").strip()
    if not raw or raw.lower() in ("none", "off"):
        return []
    return [x.strip() for x in raw.split(",") if x.strip()]


def build_grouped_lr_optimizer(model: torch.nn.Module) -> torch.optim.Optimizer:
    """Build an AdamW optimizer with separate *effective* LR for gate rows.

    Dynamic-mode gate logits live in the final router rows of
    each attention ``q_proj`` instead of in a separate Parameter. PyTorch
    optimizer groups cannot split one Parameter by row, so q_proj tensors are
    placed in the gate-lr group and :class:`RowWiseAdamW` scales their realized
    Q-row updates by ``backbone_lr / gate_lr``. Other trainable parameters use
    the backbone-lr group directly.
    """
    if getattr(model.config, "aha_mode", "dynamic") != "dynamic":
        raise ValueError("GROUPED_LR currently only supports AHA_MODE=dynamic")
    if gate_learning_rate <= 0.0 or backbone_learning_rate < 0.0:
        raise ValueError(
            "GATE_LEARNING_RATE must be positive and BACKBONE_LEARNING_RATE must be non-negative"
        )

    num_heads = model.config.num_attention_heads
    head_dim = getattr(model.config, "head_dim", model.config.hidden_size // num_heads)
    q_rows = num_heads * head_dim
    q_row_scale = backbone_learning_rate / gate_learning_rate

    gate_params = []
    gate_param_ids = set()

    row_scales = []
    for layer in model.model.layers:
        q_proj = layer.self_attn.q_proj
        for p in (q_proj.weight, q_proj.bias):
            if p is None or not p.requires_grad:
                continue
            gate_params.append(p)
            gate_param_ids.add(id(p))
            row_scales.append((p, q_rows, q_row_scale))

    backbone_params = [
        p for p in model.parameters()
        if p.requires_grad and id(p) not in gate_param_ids
    ]
    if not gate_params:
        raise ValueError("GROUPED_LR found no trainable q_proj gate parameters")

    gate_n = sum(p.numel() for p in gate_params)
    effective_gate_n = sum(
        layer.self_attn.aha_router_outputs
        * (layer.self_attn.q_proj.in_features + (layer.self_attn.q_proj.bias is not None))
        for layer in model.model.layers
    )
    backbone_n = sum(p.numel() for p in backbone_params)
    print(
        "  GROUPED_LR optimizer: "
        f"gate_lr={gate_learning_rate:.3e}, backbone_lr={backbone_learning_rate:.3e}, "
        f"q_row_update_scale={q_row_scale:.3g}, "
        f"q_proj_params={gate_n:,}, effective_gate_params={effective_gate_n:,}, "
        f"other_trainable={backbone_n:,}"
    )

    param_groups = [{"params": gate_params, "lr": gate_learning_rate}]
    if backbone_params:
        param_groups.append({"params": backbone_params, "lr": backbone_learning_rate})
    return RowWiseAdamW(
        param_groups,
        row_scales=row_scales,
        weight_decay=weight_decay,
        betas=(adam_beta1, adam_beta2),
    )


# === Configuration (override via environment variables) ===
model_path = os.environ.get("MODEL_PATH", "/workspace/AHA/models/Qwen3-0.6B")
aha_checkpoint_path = os.environ.get("AHA_CHECKPOINT_PATH", "")
dataset_path = os.environ.get("DATASET_PATH", "/workspace/Direct-Multitoken-Decoding/am-distilled-8192")
dataset_split = os.environ.get("DATASET_SPLIT", "train")
output_dir = os.environ.get("OUTPUT_DIR", "/workspace/AHA/AHA-Qwen3/ckpts/aha_qwen3_w1024")
aha_window_size = int(os.environ.get("AHA_WINDOW_SIZE", "1024"))
aha_lambda = float(os.environ.get("AHA_LAMBDA", "3e-4"))
aha_distill_weight = float(os.environ.get("AHA_DISTILL_WEIGHT", "0.0"))
# Language-modeling CE weight. Default 1.0 = standard SFT. Set to 0.0 to
# drop CE from the training objective and shape the gate purely via
# ``aux + distill`` (useful with GATE_ONLY + frozen backbone).
aha_ce_weight = float(os.environ.get("AHA_CE_WEIGHT", "1.0"))
# Hinge target on mean(gate_soft). Aux loss is ``λ · max(0, ḡ - τ)``.
# Default 1.0 keeps the legacy unconditional aux (τ=1 makes the clamp
# vacuous). Set e.g. 0.15 to cap global-attention density at 15% of
# tokens in steady state -- this is the safety floor that prevents
# gate collapse when training without CE. See AHAQwen3Config.
aha_gate_target = float(os.environ.get("AHA_GATE_TARGET", "1.0"))
# Direct Duo-style sparsity regularizer. When set to a non-negative value,
# the AHA loss uses ``AHA_REG_WEIGHT * mean(gate_soft)`` instead of the legacy
# hinge controlled by AHA_LAMBDA/AHA_GATE_TARGET.
aha_reg_weight = float(os.environ.get("AHA_REG_WEIGHT", "-1.0"))
# AHA mode: "dynamic" (per-(token, kv_head) MLP gate) or "duo" (per-(layer,
# kv_head) static scalar, DuoAttention-style). DUO mode ignores GATE_ONLY
# and instead uses its own ``DUO_ALPHA_ONLY`` freeze path.
aha_mode = os.environ.get("AHA_MODE", "dynamic")
aha_router_granularity_request = os.environ.get(
    "AHA_ROUTER_GRANULARITY", ""
).strip()
if aha_router_granularity_request and aha_router_granularity_request not in {
    "token",
    "token_kv_head",
}:
    raise ValueError(
        "AHA_ROUTER_GRANULARITY must be 'token' or 'token_kv_head', got "
        f"{aha_router_granularity_request!r}"
    )
# Override aha_local_kind on a loaded ckpt: "sink_recent" or
# "sliding_window". Useful for ablating sink dependence in dynamic mode
# from the same hot-start ckpt without retraining the source. Empty
# string keeps the ckpt's existing config value.
aha_local_kind_override = os.environ.get("AHA_LOCAL_KIND", "").strip()
duo_sink_size = int(os.environ.get("DUO_SINK_SIZE", "64"))
duo_recent_size = int(os.environ.get("DUO_RECENT_SIZE", "256"))
duo_alpha_init = float(os.environ.get("DUO_ALPHA_INIT", "1.0"))
duo_alpha_only = parse_bool("DUO_ALPHA_ONLY", default=(aha_mode == "duo"))
gate_only = parse_bool("GATE_ONLY", default=False)
grouped_lr = parse_bool("GROUPED_LR", default=False)
max_seq_length = int(os.environ.get("MAX_SEQ_LENGTH", "8192"))
min_train_tokens = get_optional_int("MIN_TRAIN_TOKENS")
max_train_tokens = get_optional_int("MAX_TRAIN_TOKENS")
max_train_samples = get_optional_int("MAX_TRAIN_SAMPLES")
learning_rate = float(os.environ.get("LEARNING_RATE", "3e-5"))
gate_learning_rate = float(os.environ.get("GATE_LEARNING_RATE", str(learning_rate)))
backbone_learning_rate = float(os.environ.get("BACKBONE_LEARNING_RATE", str(learning_rate)))
lr_scheduler_type = os.environ.get("LR_SCHEDULER_TYPE", "cosine")
warmup_ratio = float(os.environ.get("WARMUP_RATIO", "0.03"))
adam_beta1 = float(os.environ.get("ADAM_BETA1", "0.9"))
adam_beta2 = float(os.environ.get("ADAM_BETA2", "0.999"))
weight_decay = float(os.environ.get("WEIGHT_DECAY", "0.0"))
max_grad_norm = float(os.environ.get("MAX_GRAD_NORM", "1.0"))
per_device_batch_size = int(os.environ.get("PER_DEVICE_TRAIN_BATCH_SIZE", "2"))
gradient_accumulation_steps = int(os.environ.get("GRADIENT_ACCUMULATION_STEPS", "16"))
num_train_epochs = float(os.environ.get("NUM_TRAIN_EPOCHS", "1"))
max_steps = int(os.environ.get("MAX_STEPS", "-1"))  # -1 disables the cap; positive overrides num_train_epochs
logging_steps = int(os.environ.get("LOGGING_STEPS", "10"))
save_steps = int(os.environ.get("SAVE_STEPS", "500"))
save_total_limit = int(os.environ.get("SAVE_TOTAL_LIMIT", "2"))
save_strategy = os.environ.get("SAVE_STRATEGY", "steps")
seed = int(os.environ.get("SEED", "42"))
resume_from_checkpoint = resolve_resume_checkpoint(output_dir)
report_to = parse_report_to()
wandb_run_name = os.environ.get("WANDB_RUN_NAME") or os.environ.get("RUN_NAME") or os.path.basename(output_dir.rstrip("/"))
fsdp_options = parse_fsdp_options("FSDP_OPTIONS")
fsdp_transformer_layer_cls = os.environ.get("FSDP_TRANSFORMER_LAYER_CLS_TO_WRAP", "")
fsdp_activation_checkpointing = parse_bool("FSDP_ACTIVATION_CHECKPOINTING", default=True)
freeze_embeddings_lm_head = parse_bool("FREEZE_EMBEDDINGS_LM_HEAD", default=True)


def main():
    AHAQwen3Config.register_for_auto_class()
    AHAQwen3ForCausalLM.register_for_auto_class("AutoModelForCausalLM")
    os.makedirs(output_dir, exist_ok=True)

    print(f"Loading Qwen3 from {model_path} and converting to AHA-Qwen3...")
    if aha_checkpoint_path:
        print(f"  aha_checkpoint_path={aha_checkpoint_path}")
    print(
        f"  aha_mode={aha_mode}, "
        f"aha_router_granularity={aha_router_granularity_request or 'checkpoint/default'}, "
        f"aha_window_size={aha_window_size}, "
        f"aha_lambda={aha_lambda}, aha_distill_weight={aha_distill_weight}, "
        f"aha_ce_weight={aha_ce_weight}, aha_gate_target={aha_gate_target}, "
        f"aha_reg_weight={aha_reg_weight}, "
        f"gate_only={gate_only}, duo_alpha_only={duo_alpha_only}, grouped_lr={grouped_lr}"
    )
    if aha_mode == "duo":
        print(
            f"  duo_sink_size={duo_sink_size}, duo_recent_size={duo_recent_size}, "
            f"duo_alpha_init={duo_alpha_init}"
        )
    print(f"  dataset={dataset_path}[{dataset_split}]")
    print(f"  output_dir={output_dir}")
    print(f"  max_seq_length={max_seq_length}")
    print(
        "  optimizer/schedule: "
        f"lr={learning_rate:.3e}, scheduler={lr_scheduler_type}, warmup_ratio={warmup_ratio}, "
        f"betas=({adam_beta1}, {adam_beta2}), weight_decay={weight_decay}, "
        f"max_grad_norm={max_grad_norm}"
    )
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    effective_batch_size = per_device_batch_size * gradient_accumulation_steps * world_size
    print(
        "  batch: "
        f"per_device={per_device_batch_size}, grad_accum={gradient_accumulation_steps}, "
        f"world_size={world_size}, effective_global={effective_batch_size}"
    )
    if min_train_tokens is not None or max_train_tokens is not None:
        print(f"  token_filter=[{min_train_tokens}, {max_train_tokens}]")
    if max_train_samples is not None:
        print(f"  max_train_samples={max_train_samples}")
    if resume_from_checkpoint is not None:
        print(f"  resume_from_checkpoint={resume_from_checkpoint}")
    if fsdp_options:
        print(f"  fsdp_options={fsdp_options}")
        if fsdp_transformer_layer_cls:
            print(f"  fsdp_transformer_layer_cls_to_wrap={fsdp_transformer_layer_cls}")
    print(f"  report_to={report_to if report_to else 'none'}")
    if report_to and "wandb" in report_to:
        print(
            "  wandb: set WANDB_PROJECT (required for cloud UI); optional WANDB_ENTITY, WANDB_RUN_GROUP, WANDB_TAGS. "
            "Run `wandb login` once. Metrics include train/ce_loss, train/gate_aux_loss, train/loss_total, train/loss (HF)."
        )

    if aha_checkpoint_path:
        model = AHAQwen3ForCausalLM.from_pretrained_aha(
            aha_checkpoint_path,
            torch_dtype=torch.bfloat16,
            attn_implementation="sdpa",
        )
        # Config carried inside the checkpoint may predate aha_distill_weight /
        # aha_ce_weight / aha_gate_target; honor the env override either way.
        model.config.aha_distill_weight = aha_distill_weight
        model.config.aha_ce_weight = aha_ce_weight
        model.config.aha_lambda = aha_lambda
        model.config.aha_gate_target = aha_gate_target
        model.config.aha_reg_weight = aha_reg_weight
        loaded_granularity = getattr(
            model.config, "aha_router_granularity", AHA_ROUTER_GRANULARITY
        )
        if (
            aha_router_granularity_request
            and aha_router_granularity_request != loaded_granularity
        ):
            raise ValueError(
                "AHA_ROUTER_GRANULARITY does not match checkpoint architecture: "
                f"requested={aha_router_granularity_request!r}, "
                f"checkpoint={loaded_granularity!r}"
            )
        if aha_local_kind_override:
            if aha_local_kind_override not in {"sink_recent", "sliding_window"}:
                raise ValueError(
                    f"AHA_LOCAL_KIND must be 'sink_recent' or 'sliding_window', got {aha_local_kind_override!r}"
                )
            prev = getattr(model.config, "aha_local_kind", None)
            model.config.aha_local_kind = aha_local_kind_override
            print(f"  AHA_LOCAL_KIND override: {prev!r} -> {aha_local_kind_override!r}")
    else:
        model = AHAQwen3ForCausalLM.from_pretrained_qwen3(
            model_path,
            aha_window_size=aha_window_size,
            aha_lambda=aha_lambda,
            aha_distill_weight=aha_distill_weight,
            aha_ce_weight=aha_ce_weight,
            aha_gate_target=aha_gate_target,
            aha_reg_weight=aha_reg_weight,
            aha_mode=aha_mode,
            aha_router_granularity=(
                aha_router_granularity_request or AHA_ROUTER_GRANULARITY
            ),
            duo_sink_size=duo_sink_size,
            duo_recent_size=duo_recent_size,
            duo_alpha_init=duo_alpha_init,
            torch_dtype=torch.bfloat16,
            attn_implementation="sdpa",
        )

    # Historical Qwen experiments froze embeddings and lm_head. The rebuttal
    # scaling protocol disables this to match the OLMo-2 end-to-end recipe,
    # where the complete pretrained model and routers are fine-tuned jointly.
    if freeze_embeddings_lm_head:
        for param in model.model.embed_tokens.parameters():
            param.requires_grad = False
        for param in model.lm_head.parameters():
            param.requires_grad = False
    print(f"  freeze_embeddings_lm_head={freeze_embeddings_lm_head}")

    if duo_alpha_only:
        # DuoAttention-style: only the per-(layer, kv_head) scalar
        # ``full_attention_heads`` is trainable; everything else is
        # frozen. For Qwen3-0.6B this yields 224 trainable scalars.
        if getattr(model.config, "aha_mode", "dynamic") != "duo":
            raise ValueError(
                "DUO_ALPHA_ONLY=1 requires AHA_MODE=duo; otherwise there are no alpha params."
            )
        for param in model.parameters():
            param.requires_grad = False
        unfrozen = 0
        for layer in model.model.layers:
            p = layer.self_attn.full_attention_heads
            p.requires_grad = True
            unfrozen += p.numel()
        print(f"  DUO_ALPHA_ONLY: trainable scalars {unfrozen:,}")
    elif gate_only:
        # Gate-only training: freeze everything, then unfreeze just the
        # gate rows of q_proj via gradient masks. q_proj.weight / bias
        # layout is [num_heads*head_dim Q rows, native router rows].
        setup = configure_gate_only(model)
        print(
            f"  GATE_ONLY: granularity={model.config.aha_router_granularity}, "
            f"gate_rows/layer={setup.gate_rows}, "
            f"effective gate parameters {setup.effective_parameter_count:,} "
            "(Q rows grad-masked to 0)"
        )

    n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    n_total = sum(p.numel() for p in model.parameters())
    print(f"  Trainable: {n_trainable:,} / {n_total:,} ({n_trainable/n_total:.1%})")

    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

    print(f"Loading dataset from {dataset_path}...")
    dataset_dict = datasets.load_from_disk(dataset_path)
    train_dataset = dataset_dict[dataset_split]
    print(f"  loaded_rows={len(train_dataset):,}")

    if min_train_tokens is not None:
        train_dataset = train_dataset.filter(lambda example: example["num_tokens"] >= min_train_tokens)
        print(f"  rows_after_min_train_tokens={len(train_dataset):,}")
    if max_train_tokens is not None:
        train_dataset = train_dataset.filter(lambda example: example["num_tokens"] <= max_train_tokens)
        print(f"  rows_after_max_train_tokens={len(train_dataset):,}")
    if max_train_samples is not None:
        keep = min(max_train_samples, len(train_dataset))
        train_dataset = train_dataset.shuffle(seed=seed).select(range(keep))
        print(f"  rows_after_max_train_samples={len(train_dataset):,}")
    if len(train_dataset) == 0:
        raise ValueError("Training dataset is empty after filtering.")

    if int(os.environ.get("RANK", "0")) == 0:
        try:
            git_commit = subprocess.check_output(
                ["git", "rev-parse", "HEAD"], text=True
            ).strip()
            git_status = subprocess.check_output(
                ["git", "status", "--short"], text=True
            ).splitlines()
        except (OSError, subprocess.CalledProcessError):
            git_commit, git_status = None, []
        manifest = {
            "schema": "aha-qwen3-training-manifest-v1",
            "git_commit": git_commit,
            "git_status": git_status,
            "host": socket.gethostname(),
            "model_path": str(model_path),
            "source_aha_checkpoint": str(aha_checkpoint_path) or None,
            "router_granularity": getattr(
                model.config, "aha_router_granularity", AHA_ROUTER_GRANULARITY
            ),
            "native_gate_rows_per_layer": aha_router_output_size(model.config),
            "effective_gate_parameters": sum(
                layer.self_attn.aha_router_outputs
                * (
                    layer.self_attn.q_proj.in_features
                    + (layer.self_attn.q_proj.bias is not None)
                )
                for layer in model.model.layers
            ),
            "effective_sparsity_denominator": "token x KV-head x layer",
            "dataset": {
                "path": str(dataset_path),
                "split": dataset_split,
                "rows_after_filtering": len(train_dataset),
                "max_sequence_length": max_seq_length,
            },
            "training": {
                "gate_only": gate_only,
                "grouped_lr": grouped_lr,
                "learning_rate": learning_rate,
                "gate_learning_rate": gate_learning_rate,
                "backbone_learning_rate": backbone_learning_rate,
                "regularizer_weight": aha_reg_weight,
                "ce_weight": aha_ce_weight,
                "attention_distill_weight": aha_distill_weight,
                "train_gate_threshold": float(
                    os.environ.get("AHA_TRAIN_GATE_HARD_THRESHOLD", "0.5")
                ),
                "max_steps": max_steps,
                "seed": seed,
                "per_device_batch_size": per_device_batch_size,
                "gradient_accumulation_steps": gradient_accumulation_steps,
                "world_size": world_size,
            },
        }
        manifest_path = os.path.join(output_dir, "aha_training_manifest.json")
        with open(manifest_path, "w") as f:
            json.dump(manifest, f, indent=2)
    else:
        manifest_path = os.path.join(output_dir, "aha_training_manifest.json")

    fsdp_config = None
    gradient_checkpointing = True
    gradient_checkpointing_kwargs = {"use_reentrant": False}
    if fsdp_options:
        fsdp_config = {}
        if fsdp_transformer_layer_cls:
            fsdp_config["transformer_layer_cls_to_wrap"] = [fsdp_transformer_layer_cls]
        if fsdp_activation_checkpointing:
            # Transformers warns that FSDP should use fsdp_config.activation_checkpointing
            # instead of Trainer-level gradient_checkpointing to avoid redundant all-gathers.
            fsdp_config["activation_checkpointing"] = True
            gradient_checkpointing = False
            gradient_checkpointing_kwargs = None

    sft_config = SFTConfig(
        output_dir=output_dir,
        per_device_train_batch_size=per_device_batch_size,
        gradient_accumulation_steps=gradient_accumulation_steps,
        learning_rate=learning_rate,
        num_train_epochs=num_train_epochs,
        max_steps=max_steps,
        max_length=max_seq_length,
        lr_scheduler_type=lr_scheduler_type,
        warmup_ratio=warmup_ratio,
        adam_beta1=adam_beta1,
        adam_beta2=adam_beta2,
        weight_decay=weight_decay,
        max_grad_norm=max_grad_norm,
        logging_steps=logging_steps,
        save_steps=save_steps,
        save_total_limit=save_total_limit,
        save_strategy=save_strategy,
        bf16=True,
        tf32=True,
        gradient_checkpointing=gradient_checkpointing,
        gradient_checkpointing_kwargs=gradient_checkpointing_kwargs,
        dataloader_drop_last=True,
        remove_unused_columns=True,
        report_to=report_to if report_to else "none",
        run_name=wandb_run_name,
        seed=seed,
        # The custom model returns a standard local mean CE and does not use
        # ``num_items_in_batch``. Keep legacy OLMo-style per-microbatch means
        # instead of Transformers' global token-count scaling path.
        average_tokens_across_devices=False,
        fsdp=fsdp_options,
        fsdp_config=fsdp_config,
    )

    trainer = AHASFTTrainer(
        model=model,
        args=sft_config,
        train_dataset=train_dataset,
        processing_class=tokenizer,
        optimizers=(build_grouped_lr_optimizer(model), None) if grouped_lr else (None, None),
        callbacks=[AHAManifestCallback(manifest_path)],
    )

    print(
        "Starting training… Extra log fields: ce_loss, gate_aux_loss, loss_total (=ce+gate_aux), "
        "gate_soft_mean, gate_hard_mean, hf_loss_ratio (= key `loss` / loss_total). "
        "ce_loss / loss_total are raw model means (language quality + AHA aux). "
        "Key `loss` is HuggingFace Trainer’s gradient-accumulation-normalized aggregate; "
        "hf_loss_ratio should stay near 1.0."
    )
    print("Starting training...")
    train_result = trainer.train(resume_from_checkpoint=resume_from_checkpoint)
    trainer.save_model(output_dir)
    trainer.save_state()

    metrics = train_result.metrics
    trainer.log_metrics("train", metrics)
    trainer.save_metrics("train", metrics)
    print(f"Training completed. Metrics: {metrics}")


if __name__ == "__main__":
    main()