495 lines
19 KiB
Python
495 lines
19 KiB
Python
|
"""Duo-style hidden-state distillation for AHA dynamic gates.
|
||
|
|
|
||
|
|
This is the clean static-to-dynamic continuation experiment:
|
||
|
|
|
||
|
|
loss = MSE(h_full, h_dynamic) on labelled tokens + reg_weight * mean(gate_soft)
|
||
|
|
|
||
|
|
where h_full is produced by the same frozen checkpoint with every dynamic gate
|
||
|
|
forced to 1.0, and h_dynamic uses the learned per-token gate. For
|
||
|
|
``aha_mode="duo_dynamic"``, "forced to 1.0" means "reproduce the locked static
|
||
|
|
Duo full-head mask"; Duo streaming heads remain streaming. Backbone weights
|
||
|
|
stay frozen; only q_proj gate rows are trainable.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import random
|
||
|
|
import shutil
|
||
|
|
import sys
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
from torch.utils.data import DataLoader
|
||
|
|
from transformers import AutoTokenizer
|
||
|
|
|
||
|
|
HERE = os.path.dirname(os.path.abspath(__file__))
|
||
|
|
if HERE not in sys.path:
|
||
|
|
sys.path.insert(0, HERE)
|
||
|
|
|
||
|
|
from duo_train import AmDistilledDataset, LongBenchLiteAnswerDataset, collate # noqa: E402
|
||
|
|
from modeling_aha_qwen3 import ( # noqa: E402
|
||
|
|
AHA_ROUTER_GRANULARITY,
|
||
|
|
AHAQwen3Config,
|
||
|
|
AHAQwen3ForCausalLM,
|
||
|
|
aha_router_output_size,
|
||
|
|
)
|
||
|
|
from router_training_utils import configure_gate_only # noqa: E402
|
||
|
|
|
||
|
|
|
||
|
|
def _set_force_gate(model, value):
|
||
|
|
prev = getattr(model.config, "aha_force_gate_value", None)
|
||
|
|
model.config.aha_force_gate_value = value
|
||
|
|
return prev
|
||
|
|
|
||
|
|
|
||
|
|
def _restore_force_gate(model, value):
|
||
|
|
model.config.aha_force_gate_value = value
|
||
|
|
|
||
|
|
|
||
|
|
def _gate_tensors(out):
|
||
|
|
return [g.float() for g in out.all_gate_soft]
|
||
|
|
|
||
|
|
|
||
|
|
def gate_stats_from_output(out) -> dict:
|
||
|
|
with torch.no_grad():
|
||
|
|
soft = torch.cat([g.reshape(-1) for g in _gate_tensors(out)])
|
||
|
|
hard = torch.cat([g.float().reshape(-1) for g in out.all_gate_hard])
|
||
|
|
return {
|
||
|
|
"gate_soft_mean": soft.mean().item(),
|
||
|
|
"gate_soft_std": soft.std().item(),
|
||
|
|
"gate_hard_mean": hard.mean().item(),
|
||
|
|
"gate_min": soft.min().item(),
|
||
|
|
"gate_max": soft.max().item(),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def gate_param_stats(model) -> dict:
|
||
|
|
weights = []
|
||
|
|
biases = []
|
||
|
|
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
|
||
|
|
with torch.no_grad():
|
||
|
|
for layer in model.model.layers:
|
||
|
|
q_proj = layer.self_attn.q_proj
|
||
|
|
weights.append(q_proj.weight[q_rows:].detach().float().reshape(-1).cpu())
|
||
|
|
if q_proj.bias is not None:
|
||
|
|
biases.append(q_proj.bias[q_rows:].detach().float().cpu())
|
||
|
|
w = torch.cat(weights)
|
||
|
|
stats = {
|
||
|
|
"gate_weight_l2": w.norm().item(),
|
||
|
|
"gate_weight_abs_mean": w.abs().mean().item(),
|
||
|
|
}
|
||
|
|
if biases:
|
||
|
|
b = torch.cat(biases)
|
||
|
|
alpha = torch.sigmoid(b)
|
||
|
|
stats.update({
|
||
|
|
"gate_bias_mean": b.mean().item(),
|
||
|
|
"gate_bias_std": b.std().item(),
|
||
|
|
"gate_bias_alpha_mean": alpha.mean().item(),
|
||
|
|
"gate_bias_alpha_gt05": (alpha > 0.5).float().mean().item(),
|
||
|
|
})
|
||
|
|
if getattr(model.config, "aha_mode", "dynamic") == "duo_dynamic":
|
||
|
|
masks = []
|
||
|
|
with torch.no_grad():
|
||
|
|
for layer in model.model.layers:
|
||
|
|
alpha_static = layer.self_attn.full_attention_heads.detach().float()
|
||
|
|
masks.append((alpha_static > 0.5).reshape(-1).cpu())
|
||
|
|
m = torch.cat(masks).float()
|
||
|
|
stats.update({
|
||
|
|
"duo_static_full_frac": m.mean().item(),
|
||
|
|
"duo_static_streaming_frac": (1.0 - m).mean().item(),
|
||
|
|
})
|
||
|
|
return stats
|
||
|
|
|
||
|
|
|
||
|
|
def build_reg_head_weights(model, mode: str, power: float):
|
||
|
|
"""Optional per-head weights for the sparsity regularizer.
|
||
|
|
|
||
|
|
``duo_alpha_margin`` uses Duo's own static alpha as confidence: full heads
|
||
|
|
barely above 0.5 receive little sparsity pressure, while high-alpha full
|
||
|
|
heads receive normal pressure. This tests whether the alpha-low8 inference
|
||
|
|
guard can be moved into training as a smooth objective.
|
||
|
|
"""
|
||
|
|
if mode == "uniform":
|
||
|
|
return None
|
||
|
|
if mode != "duo_alpha_margin":
|
||
|
|
raise ValueError(f"unknown reg head weight mode: {mode}")
|
||
|
|
if getattr(model.config, "aha_mode", "dynamic") != "duo_dynamic":
|
||
|
|
raise ValueError(
|
||
|
|
"--reg_head_weight_mode=duo_alpha_margin requires aha_mode=duo_dynamic"
|
||
|
|
)
|
||
|
|
if power <= 0.0:
|
||
|
|
raise ValueError("--reg_head_weight_power must be positive")
|
||
|
|
|
||
|
|
weights = []
|
||
|
|
with torch.no_grad():
|
||
|
|
for layer in model.model.layers:
|
||
|
|
alpha = layer.self_attn.full_attention_heads.detach().float().clamp(0.0, 1.0)
|
||
|
|
duo_full = (alpha > 0.5).float()
|
||
|
|
margin = ((alpha - 0.5) / 0.5).clamp(0.0, 1.0).pow(power)
|
||
|
|
weights.append((margin * duo_full).view(1, 1, -1))
|
||
|
|
flat = torch.cat([w.reshape(-1).cpu() for w in weights])
|
||
|
|
print(
|
||
|
|
"[dynamic-duo] reg_head_weight_mode=duo_alpha_margin "
|
||
|
|
f"power={power:g} nonzero={int((flat > 0).sum().item())}/{flat.numel()} "
|
||
|
|
f"mean={flat.mean().item():.4f}",
|
||
|
|
flush=True,
|
||
|
|
)
|
||
|
|
return weights
|
||
|
|
|
||
|
|
|
||
|
|
def sparsity_reg_from_gates(out, reg_head_weights):
|
||
|
|
gate_layers = _gate_tensors(out)
|
||
|
|
if reg_head_weights is None:
|
||
|
|
gate_soft = torch.cat([g.reshape(-1) for g in gate_layers])
|
||
|
|
return gate_soft.mean()
|
||
|
|
|
||
|
|
weighted = []
|
||
|
|
for gate, weight in zip(gate_layers, reg_head_weights):
|
||
|
|
w = weight.to(device=gate.device, dtype=gate.dtype)
|
||
|
|
weighted.append((gate * w).reshape(-1))
|
||
|
|
if not weighted:
|
||
|
|
raise ValueError("empty weighted sparsity regularizer")
|
||
|
|
return torch.cat(weighted).mean()
|
||
|
|
|
||
|
|
|
||
|
|
def build_gate_optimizer(model, lr: float):
|
||
|
|
setup = configure_gate_only(model)
|
||
|
|
print(
|
||
|
|
f"[dynamic-duo] router_granularity={model.config.aha_router_granularity} "
|
||
|
|
f"gate_rows/layer={setup.gate_rows} effective trainable gate params: "
|
||
|
|
f"{setup.effective_parameter_count:,}",
|
||
|
|
flush=True,
|
||
|
|
)
|
||
|
|
return torch.optim.AdamW([{"params": setup.parameters, "lr": lr}], weight_decay=0.0)
|
||
|
|
|
||
|
|
|
||
|
|
def save_checkpoint(model, tokenizer, output_dir: str, step: int, stats: dict):
|
||
|
|
sub = os.path.join(output_dir, f"checkpoint-{step}")
|
||
|
|
os.makedirs(sub, exist_ok=True)
|
||
|
|
prev_force = getattr(model.config, "aha_force_gate_value", None)
|
||
|
|
model.config.aha_force_gate_value = None
|
||
|
|
model.save_pretrained(sub, safe_serialization=True)
|
||
|
|
tokenizer.save_pretrained(sub)
|
||
|
|
model.config.aha_force_gate_value = prev_force
|
||
|
|
with open(os.path.join(sub, "dynamic_duo_state.json"), "w") as f:
|
||
|
|
json.dump(
|
||
|
|
{
|
||
|
|
"step": step,
|
||
|
|
"router_granularity": getattr(
|
||
|
|
model.config,
|
||
|
|
"aha_router_granularity",
|
||
|
|
AHA_ROUTER_GRANULARITY,
|
||
|
|
),
|
||
|
|
"native_gate_rows_per_layer": aha_router_output_size(model.config),
|
||
|
|
"effective_sparsity_denominator": "token x KV-head x layer",
|
||
|
|
**stats,
|
||
|
|
**gate_param_stats(model),
|
||
|
|
},
|
||
|
|
f,
|
||
|
|
indent=2,
|
||
|
|
)
|
||
|
|
args_manifest = os.path.join(output_dir, "dynamic_duo_train_args.json")
|
||
|
|
if os.path.exists(args_manifest):
|
||
|
|
shutil.copyfile(
|
||
|
|
args_manifest,
|
||
|
|
os.path.join(sub, "dynamic_duo_train_args.json"),
|
||
|
|
)
|
||
|
|
print(f"[dynamic-duo] saved {sub}", flush=True)
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
AHAQwen3Config.register_for_auto_class()
|
||
|
|
AHAQwen3ForCausalLM.register_for_auto_class("AutoModelForCausalLM")
|
||
|
|
p = argparse.ArgumentParser()
|
||
|
|
p.add_argument("--aha_checkpoint", required=True)
|
||
|
|
p.add_argument("--model_path", default="/workspace/AHA/models/Qwen3-0.6B")
|
||
|
|
p.add_argument("--output_dir", required=True)
|
||
|
|
p.add_argument("--am_dataset_path", default="/workspace/Direct-Multitoken-Decoding/am-distilled-8192")
|
||
|
|
p.add_argument("--am_dataset_split", default="train")
|
||
|
|
p.add_argument(
|
||
|
|
"--data_source",
|
||
|
|
default="am_distilled",
|
||
|
|
choices=["am_distilled", "longbench_lite"],
|
||
|
|
help=(
|
||
|
|
"Training data for dynamic gate continuation. The default keeps the "
|
||
|
|
"legacy AM-distilled behavior. longbench_lite is a diagnostic "
|
||
|
|
"target-distribution calibration source, not a paper main protocol."
|
||
|
|
),
|
||
|
|
)
|
||
|
|
p.add_argument(
|
||
|
|
"--am_label_mode",
|
||
|
|
default="full",
|
||
|
|
choices=["full", "answer_only"],
|
||
|
|
help=(
|
||
|
|
"Label mask for --data_source=am_distilled. full keeps all-token "
|
||
|
|
"hidden-state distill; answer_only distills only the final answer span."
|
||
|
|
),
|
||
|
|
)
|
||
|
|
p.add_argument(
|
||
|
|
"--longbench_tasks",
|
||
|
|
nargs="+",
|
||
|
|
default=["passage_retrieval_en", "multifieldqa_en", "qasper", "2wikimqa"],
|
||
|
|
help="Tasks used when --data_source=longbench_lite.",
|
||
|
|
)
|
||
|
|
p.add_argument("--longbench_samples_per_task", type=int, default=30)
|
||
|
|
p.add_argument("--longbench_cache_dir", default="/workspace/AHA/AHA-Qwen3/data/longbench_cache")
|
||
|
|
p.add_argument("--max_length", type=int, default=8192)
|
||
|
|
p.add_argument("--num_steps", type=int, default=400)
|
||
|
|
p.add_argument("--warmup_ratio", type=float, default=0.2)
|
||
|
|
p.add_argument("--lr", type=float, default=3e-5)
|
||
|
|
p.add_argument("--reg_weight", type=float, default=0.05)
|
||
|
|
p.add_argument(
|
||
|
|
"--reg_head_weight_mode",
|
||
|
|
default="uniform",
|
||
|
|
choices=["uniform", "duo_alpha_margin"],
|
||
|
|
help=(
|
||
|
|
"Per-head weighting for the sparsity regularizer. uniform keeps "
|
||
|
|
"the original mean(gate_soft). duo_alpha_margin downweights Duo "
|
||
|
|
"static-full heads close to alpha=0.5 so fragile boundary full "
|
||
|
|
"heads are not pushed local as strongly."
|
||
|
|
),
|
||
|
|
)
|
||
|
|
p.add_argument(
|
||
|
|
"--reg_head_weight_power",
|
||
|
|
type=float,
|
||
|
|
default=1.0,
|
||
|
|
help=(
|
||
|
|
"Power applied to the Duo alpha margin when "
|
||
|
|
"--reg_head_weight_mode=duo_alpha_margin."
|
||
|
|
),
|
||
|
|
)
|
||
|
|
p.add_argument("--ce_weight", type=float, default=0.0)
|
||
|
|
p.add_argument(
|
||
|
|
"--distill_tail_frac",
|
||
|
|
type=float,
|
||
|
|
default=0.0,
|
||
|
|
help=(
|
||
|
|
"If >0, add a tail-aware hidden-state distill term over the top "
|
||
|
|
"fraction of labelled-token MSE values. This keeps the standard "
|
||
|
|
"mean distill objective but prevents rare long-retrieval errors "
|
||
|
|
"from being averaged away."
|
||
|
|
),
|
||
|
|
)
|
||
|
|
p.add_argument(
|
||
|
|
"--distill_tail_weight",
|
||
|
|
type=float,
|
||
|
|
default=0.0,
|
||
|
|
help="Weight for the top-token MSE term enabled by --distill_tail_frac.",
|
||
|
|
)
|
||
|
|
p.add_argument("--batch_size", type=int, default=1)
|
||
|
|
p.add_argument("--grad_accum", type=int, default=1)
|
||
|
|
p.add_argument("--save_steps", type=int, default=100)
|
||
|
|
p.add_argument("--log_steps", type=int, default=10)
|
||
|
|
p.add_argument("--seed", type=int, default=42)
|
||
|
|
p.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"])
|
||
|
|
p.add_argument("--attn_impl", default="sdpa", choices=["sdpa", "eager"])
|
||
|
|
p.add_argument("--aha_local_kind", default="sink_recent", choices=["sink_recent", "sliding_window"])
|
||
|
|
p.add_argument(
|
||
|
|
"--router_granularity",
|
||
|
|
default=None,
|
||
|
|
choices=["token", "token_kv_head"],
|
||
|
|
help="Assert that the input checkpoint has this persisted router granularity.",
|
||
|
|
)
|
||
|
|
args = p.parse_args()
|
||
|
|
|
||
|
|
torch.manual_seed(args.seed)
|
||
|
|
random.seed(args.seed)
|
||
|
|
np.random.seed(args.seed)
|
||
|
|
os.makedirs(args.output_dir, exist_ok=True)
|
||
|
|
with open(os.path.join(args.output_dir, "dynamic_duo_train_args.json"), "w") as f:
|
||
|
|
json.dump(vars(args), f, indent=2)
|
||
|
|
|
||
|
|
dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[args.dtype]
|
||
|
|
print(f"[dynamic-duo] checkpoint={args.aha_checkpoint}", flush=True)
|
||
|
|
print(
|
||
|
|
f"[dynamic-duo] loss=hidden_state_distill + {args.reg_weight} * mean(gate_soft), "
|
||
|
|
f"ce_weight={args.ce_weight} tail_frac={args.distill_tail_frac} "
|
||
|
|
f"tail_weight={args.distill_tail_weight}",
|
||
|
|
flush=True,
|
||
|
|
)
|
||
|
|
if not (0.0 <= args.distill_tail_frac <= 1.0):
|
||
|
|
raise ValueError("--distill_tail_frac must be in [0, 1]")
|
||
|
|
if args.distill_tail_weight < 0.0:
|
||
|
|
raise ValueError("--distill_tail_weight must be non-negative")
|
||
|
|
|
||
|
|
model = AHAQwen3ForCausalLM.from_pretrained_aha(
|
||
|
|
args.aha_checkpoint,
|
||
|
|
torch_dtype=dtype,
|
||
|
|
attn_implementation=args.attn_impl,
|
||
|
|
).cuda()
|
||
|
|
if getattr(model.config, "aha_mode", "dynamic") not in ("dynamic", "duo_dynamic"):
|
||
|
|
raise ValueError("dynamic_duo_train.py requires a dynamic or duo_dynamic AHA checkpoint")
|
||
|
|
loaded_granularity = getattr(
|
||
|
|
model.config, "aha_router_granularity", AHA_ROUTER_GRANULARITY
|
||
|
|
)
|
||
|
|
if args.router_granularity and args.router_granularity != loaded_granularity:
|
||
|
|
raise ValueError(
|
||
|
|
"--router_granularity does not match checkpoint architecture: "
|
||
|
|
f"requested={args.router_granularity!r}, checkpoint={loaded_granularity!r}"
|
||
|
|
)
|
||
|
|
model.config.aha_local_kind = args.aha_local_kind
|
||
|
|
model.config.aha_distill_weight = 0.0
|
||
|
|
model.config.aha_ce_weight = 0.0
|
||
|
|
model.config.aha_reg_weight = -1.0
|
||
|
|
model.config.aha_force_gate_value = None
|
||
|
|
print(
|
||
|
|
f"[dynamic-duo] router_granularity={loaded_granularity} "
|
||
|
|
f"native_gate_rows={aha_router_output_size(model.config)} "
|
||
|
|
f"local_kind={model.config.aha_local_kind} "
|
||
|
|
f"sink={getattr(model.config, 'duo_sink_size', None)} "
|
||
|
|
f"recent={getattr(model.config, 'duo_recent_size', None)}",
|
||
|
|
flush=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
|
||
|
|
if tokenizer.pad_token_id is None:
|
||
|
|
tokenizer.pad_token = tokenizer.eos_token
|
||
|
|
|
||
|
|
optim = build_gate_optimizer(model, args.lr)
|
||
|
|
reg_head_weights = build_reg_head_weights(
|
||
|
|
model, args.reg_head_weight_mode, args.reg_head_weight_power
|
||
|
|
)
|
||
|
|
model.enable_input_require_grads()
|
||
|
|
model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
|
||
|
|
|
||
|
|
if args.data_source == "am_distilled":
|
||
|
|
dataset = AmDistilledDataset(
|
||
|
|
ds_path=args.am_dataset_path,
|
||
|
|
split=args.am_dataset_split,
|
||
|
|
max_length=args.max_length,
|
||
|
|
seed=args.seed,
|
||
|
|
tokenizer=tokenizer,
|
||
|
|
label_mode=args.am_label_mode,
|
||
|
|
)
|
||
|
|
print(
|
||
|
|
f"[dynamic-duo] data_source=am_distilled path={args.am_dataset_path} "
|
||
|
|
f"split={args.am_dataset_split} n={len(dataset):,} "
|
||
|
|
f"label_mode={args.am_label_mode} max_length={args.max_length}",
|
||
|
|
flush=True,
|
||
|
|
)
|
||
|
|
elif args.data_source == "longbench_lite":
|
||
|
|
dataset = LongBenchLiteAnswerDataset(
|
||
|
|
tokenizer=tokenizer,
|
||
|
|
tasks=args.longbench_tasks,
|
||
|
|
samples_per_task=args.longbench_samples_per_task,
|
||
|
|
max_length=args.max_length,
|
||
|
|
seed=args.seed,
|
||
|
|
cache_dir=args.longbench_cache_dir,
|
||
|
|
)
|
||
|
|
print(
|
||
|
|
f"[dynamic-duo] data_source=longbench_lite tasks={args.longbench_tasks} "
|
||
|
|
f"samples_per_task={args.longbench_samples_per_task} "
|
||
|
|
f"rows={len(dataset.rows):,} max_length={args.max_length}",
|
||
|
|
flush=True,
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
raise ValueError(f"unknown data_source: {args.data_source}")
|
||
|
|
loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, collate_fn=collate, num_workers=0)
|
||
|
|
data_iter = iter(loader)
|
||
|
|
warm = max(1, int(args.num_steps * args.warmup_ratio))
|
||
|
|
|
||
|
|
def lr_at(step):
|
||
|
|
if step < warm:
|
||
|
|
return max(0.1, (step + 1) / warm)
|
||
|
|
if step > args.num_steps - warm:
|
||
|
|
return max(0.1, (args.num_steps - step) / warm)
|
||
|
|
return 1.0
|
||
|
|
|
||
|
|
model.train()
|
||
|
|
running = {"distill": 0.0, "reg": 0.0, "ce": 0.0, "loss": 0.0, "gate_soft": 0.0, "gate_hard": 0.0}
|
||
|
|
steps_in_window = 0
|
||
|
|
|
||
|
|
for step in range(args.num_steps):
|
||
|
|
try:
|
||
|
|
batch = next(data_iter)
|
||
|
|
except StopIteration:
|
||
|
|
data_iter = iter(loader)
|
||
|
|
batch = next(data_iter)
|
||
|
|
|
||
|
|
input_ids = batch["input_ids"].cuda()
|
||
|
|
labels = batch["labels"].cuda()
|
||
|
|
label_mask = labels != -100
|
||
|
|
|
||
|
|
prev_force = _set_force_gate(model, 1.0)
|
||
|
|
with torch.no_grad():
|
||
|
|
out_full = model.model(input_ids=input_ids, use_cache=False)
|
||
|
|
h_full = out_full.last_hidden_state
|
||
|
|
_restore_force_gate(model, prev_force)
|
||
|
|
|
||
|
|
out_mix = model.model(input_ids=input_ids, use_cache=False)
|
||
|
|
h_mix = out_mix.last_hidden_state
|
||
|
|
|
||
|
|
if label_mask.any():
|
||
|
|
diff = (h_full.float() - h_mix.float())[label_mask]
|
||
|
|
else:
|
||
|
|
diff = (h_full.float() - h_mix.float()).reshape(-1, h_mix.shape[-1])
|
||
|
|
token_mse = diff.pow(2).mean(dim=-1)
|
||
|
|
distill_mean = token_mse.mean()
|
||
|
|
if args.distill_tail_frac > 0.0 and args.distill_tail_weight > 0.0:
|
||
|
|
k = max(1, int(np.ceil(token_mse.numel() * args.distill_tail_frac)))
|
||
|
|
distill_tail = torch.topk(token_mse, k=k, largest=True).values.mean()
|
||
|
|
distill = distill_mean + args.distill_tail_weight * distill_tail
|
||
|
|
else:
|
||
|
|
distill_tail = h_mix.new_zeros((), dtype=torch.float32)
|
||
|
|
distill = distill_mean
|
||
|
|
|
||
|
|
reg = sparsity_reg_from_gates(out_mix, reg_head_weights)
|
||
|
|
|
||
|
|
if args.ce_weight > 0.0:
|
||
|
|
logits = model.lm_head(h_mix).float()
|
||
|
|
shift_logits = logits[:, :-1, :].contiguous()
|
||
|
|
shift_labels = labels[:, 1:].contiguous()
|
||
|
|
ce = torch.nn.functional.cross_entropy(
|
||
|
|
shift_logits.view(-1, shift_logits.size(-1)),
|
||
|
|
shift_labels.view(-1),
|
||
|
|
ignore_index=-100,
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
ce = h_mix.new_zeros((), dtype=torch.float32)
|
||
|
|
|
||
|
|
loss = distill + args.reg_weight * reg + args.ce_weight * ce
|
||
|
|
(loss / args.grad_accum).backward()
|
||
|
|
|
||
|
|
if (step + 1) % args.grad_accum == 0:
|
||
|
|
for group in optim.param_groups:
|
||
|
|
group["lr"] = args.lr * lr_at(step)
|
||
|
|
optim.step()
|
||
|
|
optim.zero_grad()
|
||
|
|
|
||
|
|
stats = gate_stats_from_output(out_mix)
|
||
|
|
running["distill"] += float(distill.detach())
|
||
|
|
running["reg"] += float(reg.detach())
|
||
|
|
running["ce"] += float(ce.detach())
|
||
|
|
running["loss"] += float(loss.detach())
|
||
|
|
running["gate_soft"] += stats["gate_soft_mean"]
|
||
|
|
running["gate_hard"] += stats["gate_hard_mean"]
|
||
|
|
steps_in_window += 1
|
||
|
|
|
||
|
|
if (step + 1) % args.log_steps == 0:
|
||
|
|
denom = float(steps_in_window)
|
||
|
|
print(
|
||
|
|
f"[step {step + 1:4d}/{args.num_steps}] "
|
||
|
|
f"distill={running['distill']/denom:.6f} "
|
||
|
|
f"reg={running['reg']/denom:.6f} "
|
||
|
|
f"loss={running['loss']/denom:.6f} "
|
||
|
|
f"gate_soft={running['gate_soft']/denom:.4f} "
|
||
|
|
f"gate_hard={running['gate_hard']/denom:.4f} "
|
||
|
|
f"lr={optim.param_groups[0]['lr']:.4e} "
|
||
|
|
f"seq_len={input_ids.shape[1]}",
|
||
|
|
flush=True,
|
||
|
|
)
|
||
|
|
running = {k: 0.0 for k in running}
|
||
|
|
steps_in_window = 0
|
||
|
|
|
||
|
|
if (step + 1) % args.save_steps == 0 or (step + 1) == args.num_steps:
|
||
|
|
save_checkpoint(model, tokenizer, args.output_dir, step + 1, stats)
|
||
|
|
|
||
|
|
print(f"[dynamic-duo] done. Final gate param stats: {gate_param_stats(model)}", flush=True)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|