643 lines
29 KiB
Python
643 lines
29 KiB
Python
"""
|
||
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()
|