| #!/usr/bin/env python3 | |
| # ------------------------------------------------------------------------ | |
| # Copyright (c) 2024-present, BAAI. All Rights Reserved. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # ------------------------------------------------------------------------ | |
| """URSA one-step distillation trainer (DiMO-style), 8-GPU distributed. | |
| Verified native inference regime (from A/B testing — ground truth): | |
| height=320, width=512, num_frames=49, guidance_scale=7, teacher_steps=50. | |
| no_cfg (guidance_scale=1) is NOT a valid baseline for this URSA checkpoint. | |
| Defaults in configs/distill_dimo.yaml are aligned to this regime. | |
| Launch command: | |
| accelerate launch --config_file accelerate_configs/deepspeed_zero2.yaml \\ | |
| --machine_rank 0 --num_machines 1 --num_processes 8 \\ | |
| scripts/train_distill_dimo.py \\ | |
| config="./configs/distill_dimo.yaml" \\ | |
| experiment.output_dir="./experiments/distill_dimo" \\ | |
| distill.teacher_ckpt="/path/to/URSA-1.7B-IBQ1024" \\ | |
| distill.prompt_source="/data/Koala_36M_*.csv" \\ | |
| distill.batch_size_per_gpu=1 | |
| Smoke test (single-GPU, 50 steps): | |
| accelerate launch --num_processes 1 \\ | |
| scripts/train_distill_dimo.py \\ | |
| config="./configs/distill_dimo.yaml" \\ | |
| experiment.output_dir="./experiments/smoke" \\ | |
| distill.teacher_ckpt="/path/to/URSA-1.7B-IBQ1024" \\ | |
| distill.prompt_source="prompts.txt" \\ | |
| training.max_train_steps=50 | |
| Algorithm summary (9 stages per iteration) | |
| ------------------------------------------ | |
| Stage 1 Tokenize → txt_ids [B, L] (CPU in worker, moved to GPU in run_step) | |
| Stage 2 x_init ~ Uniform(K) (+ p_init mixing from x_hat_prev) | |
| Stage 3 no_grad student(x_init) → x_hat [B, N], logp for PG | |
| Stage 4 x_t = scheduler.add_noise(x_hat_4d, t) [B,T,H,W], long | |
| Stage 5 no_grad teacher(x_t) → z_T_cond [B,N,K] (+ uncond if CFG) | |
| Stage 6 aux update × fake_rounds: Jeffrey(z_T_target, z_A_cond).backward() | |
| Stage 7 student KD forward on x_t → z_S_cond [B,N,K] | |
| Stage 8 reward = -KL(z_T_cond, z_S_cond) [detached]; adv = reward - baseline_ema | |
| Stage 9 Two-backward: | |
| 9a _no_sync_backward(lambda_kd * loss_kd) [frees KD graph] | |
| 9b accelerator.backward(lambda_pg * loss_pg - lambda_ent * H_mean) | |
| opt_student.step() | |
| """ | |
| import collections | |
| import copy | |
| import os | |
| import sys | |
| from typing import Optional | |
| import torch | |
| import torch.nn.functional as F | |
| from torch.utils.data import DataLoader | |
| _ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) | |
| if _ROOT not in sys.path: | |
| sys.path.insert(0, _ROOT) | |
| from diffnext.engine import engine_utils | |
| from diffnext.engine.lr_scheduler import CosineLR | |
| from diffnext.pipelines.ursa.pipeline_ursa_distill_dimo import ( | |
| URSADistillDiMOPipeline, | |
| _get_logits, | |
| _stable_kl, | |
| _stable_jeffrey, | |
| _build_guided_logits, | |
| _cfg_warmup_prob, | |
| _no_sync_backward, | |
| _reset_flex_attn, | |
| VERIFIED_NATIVE_DEFAULTS, | |
| check_verified_regime, | |
| ) | |
| from diffnext.utils import accelerate_utils | |
| from diffnext.utils import omegaconf_utils | |
| from diffnext.utils import profiler | |
| from src.distill.prompt_dataset import ( | |
| CSVSpec, | |
| InfiniteDataLoader, | |
| PromptDataset, | |
| make_collate_fn, | |
| ) | |
| # --------------------------------------------------------------------------- | |
| # DistillTwinModel — single nn.Module wrapping student + aux for DeepSpeed | |
| # --------------------------------------------------------------------------- | |
| class DistillTwinModel(torch.nn.Module): | |
| """Wrapper that holds both student and aux as sub-modules. | |
| DeepSpeed (via Accelerate) only allows a single model in | |
| ``accelerator.prepare()``. This container satisfies that constraint | |
| while keeping student and aux as separately addressable sub-modules | |
| with independent param groups. | |
| """ | |
| def __init__(self, student: torch.nn.Module, aux: torch.nn.Module): | |
| super().__init__() | |
| self.student = student | |
| self.aux = aux | |
| def forward(self, which: str, input_ids, rope_pos=None, **kwargs): | |
| if which == "student": | |
| return self.student(input_ids, rope_pos=rope_pos, **kwargs) | |
| elif which == "aux": | |
| return self.aux(input_ids, rope_pos=rope_pos, **kwargs) | |
| else: | |
| raise ValueError(f"DistillTwinModel: unknown sub-model '{which}'") | |
| # --------------------------------------------------------------------------- | |
| # DistillTrainer | |
| # --------------------------------------------------------------------------- | |
| class DistillTrainer: | |
| """Training orchestrator for on-policy one-step distillation. | |
| Reuses the same accelerate / logger / checkpoint API as | |
| ``diffnext.engine.train_engine.Trainer`` so the distributed setup is | |
| identical to the original training framework. | |
| Key differences from standard Trainer: | |
| - Three models (teacher frozen, student + aux trainable) | |
| - Student and aux are wrapped in a single ``DistillTwinModel`` so that | |
| only one ``accelerator.prepare()`` call is needed (DeepSpeed requirement) | |
| - One optimizer with two param_groups: [0]=student, [1]=aux | |
| - LR schedulers for both param groups | |
| - Two-backward strategy within each step | |
| - PromptDataset (no video latents; prompt-only) | |
| - Stage 6 freezes student / unfreezes aux; Stages 7-9 do the reverse | |
| """ | |
| def __init__(self, config, accelerator, logger): | |
| self.config = config | |
| self.accelerator = accelerator | |
| self.logger = logger | |
| cfg = config.distill | |
| dtype = accelerate_utils.precision_to_dtype(config.training.mixed_precision) | |
| self.device = accelerator.device | |
| # -------- Pipeline (teacher + student + aux) ---------------------- | |
| logger.info(f"[init] Loading teacher from {cfg.teacher_ckpt} ...") | |
| self.pipe = URSADistillDiMOPipeline( | |
| teacher_ckpt=cfg.teacher_ckpt, | |
| compute_dtype=dtype, | |
| aux_noise_std=float(cfg.get("aux_noise_std", 0.0)), | |
| ) | |
| # Move teacher to GPU (not prepared by accelerate — frozen). | |
| self.pipe.teacher = self.pipe.teacher.to(self.device) | |
| self.pipe.scheduler.to(device=self.device) | |
| # Compute latents shape from video geometry. | |
| from src.distill.utils_ursa_inputs import compute_latents_shape | |
| # Read VAE strides from pipeline (falls back to URSA defaults 4/8). | |
| vae_t = int(getattr(self.pipe, "vae_temporal_stride", 4)) | |
| vae_s = int(getattr(self.pipe, "vae_spatial_stride", 8)) | |
| self.latents_shape = compute_latents_shape( | |
| cfg.num_frames, cfg.height, cfg.width, vae_t, vae_s | |
| ) | |
| T, H, W = self.latents_shape | |
| self.N = T * H * W | |
| self.K = self.pipe.codebook_size | |
| logger.info( | |
| f"[init] latents_shape=({T},{H},{W}) N={self.N} K={self.K} " | |
| f"CFG={'ON' if cfg.enable_teacher_cfg else 'OFF'}" | |
| ) | |
| # Pre-compute uncond token IDs (empty string, [1, L]) on CPU. | |
| self.txt_uncond_base_cpu = self.pipe.tokenizer( | |
| [""], | |
| max_length=int(cfg.max_prompt_length), | |
| padding="max_length", | |
| padding_side="left", | |
| truncation=True, | |
| return_tensors="pt", | |
| ).input_ids # [1, L] CPU | |
| # -------- Optimizers (before accelerate.prepare) ------------------ | |
| # Single optimizer with two param groups: | |
| # group[0] = student params, group[1] = aux params | |
| opt_cls = torch.optim.AdamW | |
| opt_s_params = dict( | |
| lr=float(config.optimizer_student.params.lr), | |
| betas=tuple(config.optimizer_student.params.get("betas", [0.9, 0.95])), | |
| weight_decay=float(config.optimizer_student.params.get("weight_decay", 0.01)), | |
| ) | |
| opt_a_params = dict( | |
| lr=float(config.optimizer_aux.params.lr), | |
| betas=tuple(config.optimizer_aux.params.get("betas", [0.9, 0.95])), | |
| weight_decay=float(config.optimizer_aux.params.get("weight_decay", 0.01)), | |
| ) | |
| def _enable_gcpt(m): | |
| # m.model.layers 是 Qwen3Model 的层列表 | |
| for layer in m.model.layers: | |
| layer.gradient_checkpointing = True | |
| layer.self_attn.gradient_checkpointing = True | |
| layer.mlp.gradient_checkpointing = True | |
| _enable_gcpt(self.pipe.student) | |
| _enable_gcpt(self.pipe.aux) | |
| # -------- 断点续传：在 ZeRO-3 切分参数前加载权重 ------------------ | |
| self.global_step = int(config.experiment.get("resume_iter", 0)) | |
| if self.global_step > 0: | |
| ckpt_dir = os.path.join( | |
| config.experiment.output_dir, "checkpoints", f"checkpoint-{self.global_step}" | |
| ) | |
| if os.path.exists(ckpt_dir): | |
| logger.info(f"[Resume] 正在从 {ckpt_dir} 恢复 Student 和 Aux 的权重...") | |
| # 必须在 map_location="cpu" 下加载，防止爆显存，随后 prepare 会自动分配 | |
| self.pipe.student.load_state_dict(torch.load(os.path.join(ckpt_dir, "student.pt"), map_location="cpu")) | |
| self.pipe.aux.load_state_dict(torch.load(os.path.join(ckpt_dir, "aux.pt"), map_location="cpu")) | |
| else: | |
| logger.warning(f"[Resume] 找不到检查点 {ckpt_dir}，将从随机初始状态起步！") | |
| # -------- Wrap student + aux into a single DistillTwinModel -------- | |
| twin_model = DistillTwinModel(self.pipe.student, self.pipe.aux) | |
| opt_raw = opt_cls([ | |
| {"params": list(self.pipe.student.parameters()), **opt_s_params}, | |
| {"params": list(self.pipe.aux.parameters()), **opt_a_params}, | |
| ]) | |
| # -------- accelerate.prepare: single model + single optimizer ------ | |
| # Teacher is NOT prepared (frozen; no grad sync needed). | |
| self.model, self.optimizer = accelerator.prepare(twin_model, opt_raw) | |
| # LR schedulers (step() called manually at end of each step). | |
| self.scheduler_s = CosineLR( | |
| lr_max=float(config.optimizer_student.params.lr), | |
| lr_min=float(config.lr_scheduler.params.get("lr_min", 1e-6)), | |
| max_steps=int(config.training.max_train_steps), | |
| warmup_steps=int(config.lr_scheduler.params.get("warmup_steps", 500)), | |
| ) | |
| self.scheduler_a = CosineLR( | |
| lr_max=float(config.optimizer_aux.params.lr), | |
| lr_min=float(config.lr_scheduler.params.get("lr_min", 1e-6)), | |
| max_steps=int(config.training.max_train_steps), | |
| warmup_steps=int(config.lr_scheduler.params.get("warmup_steps", 500)), | |
| ) | |
| # -------- Dataset / DataLoader ------------------------------------ | |
| dataloader_cfg = config.get("prompt_dataloader", {}) | |
| dataset = PromptDataset( | |
| prompt_source=str(cfg.prompt_source), | |
| shuffle_files=bool(dataloader_cfg.get("shuffle_files", True)), | |
| shuffle_buffer=int(dataloader_cfg.get("shuffle_buffer", 0)), | |
| seed=int(config.training.seed), | |
| infinite=True, | |
| csv=CSVSpec(caption_field=str(dataloader_cfg.get("caption_field", "caption"))), | |
| ) | |
| # collate_fn: tokenize on CPU (no CUDA in workers). | |
| collate_fn = make_collate_fn( | |
| self.pipe.tokenizer, | |
| max_prompt_length=int(cfg.max_prompt_length), | |
| device=torch.device("cpu"), # CPU output — moved to GPU in run_step | |
| ) | |
| loader = DataLoader( | |
| dataset, | |
| batch_size=int(cfg.batch_size_per_gpu), | |
| shuffle=False, # IterableDataset: no shuffle flag | |
| drop_last=True, | |
| num_workers=int(dataloader_cfg.get("num_workers", 2)), | |
| collate_fn=collate_fn, | |
| pin_memory=True, | |
| ) | |
| # DataLoader is NOT prepared by accelerate because PromptDataset | |
| # handles per-rank file sharding internally via torch.distributed. | |
| self._inf_loader = InfiniteDataLoader(loader) | |
| # -------- Training state ------------------------------------------ | |
| # self.global_step = int(config.experiment.get("resume_iter", 0)) | |
| self.baseline_ema: float = 0.0 | |
| self.x_hat_prev: Optional[torch.Tensor] = None | |
| self.metrics = collections.OrderedDict() | |
| # -------- Verified regime validation -------------------------------- | |
| native = VERIFIED_NATIVE_DEFAULTS | |
| is_native = check_verified_regime( | |
| height=int(cfg.height), | |
| width=int(cfg.width), | |
| num_frames=int(cfg.num_frames), | |
| guidance_scale=float(cfg.teacher_cfg_scale) if cfg.enable_teacher_cfg else None, | |
| label="train", | |
| ) | |
| logger.info( | |
| f"[init] verified_native_regime={is_native} " | |
| f"geometry=({cfg.num_frames}×{cfg.height}×{cfg.width}) " | |
| f"teacher_cfg_scale={cfg.teacher_cfg_scale if cfg.enable_teacher_cfg else 'OFF'}" | |
| ) | |
| if not cfg.enable_teacher_cfg: | |
| logger.warning( | |
| "[WARN] Teacher CFG is DISABLED. no_cfg is known to produce " | |
| "blank/blurry output for this URSA checkpoint. " | |
| "Distillation without CFG is unlikely to produce useful results." | |
| ) | |
| elif float(cfg.teacher_cfg_scale) != native["guidance_scale"]: | |
| logger.warning( | |
| f"[WARN] teacher_cfg_scale={cfg.teacher_cfg_scale} differs from " | |
| f"the verified working value ({native['guidance_scale']}). " | |
| "Outputs may deviate from the official inference working point." | |
| ) | |
| logger.info( | |
| f"[init] student params: {engine_utils.count_params(self.pipe.student):.2f}M" | |
| ) | |
| logger.info( | |
| f"[init] max_train_steps={config.training.max_train_steps} " | |
| f"batch_size_per_gpu={cfg.batch_size_per_gpu} " | |
| f"num_processes={accelerator.num_processes}" | |
| ) | |
| # ----------------------------------------------------------------------- | |
| # run_step: Stages 1-9 | |
| # ----------------------------------------------------------------------- | |
| def run_step(self, step: int) -> dict: | |
| """Execute one distillation step (Stages 1-9).""" | |
| cfg = self.config.distill | |
| T, H, W = self.latents_shape | |
| N, K = self.N, self.K | |
| device = self.device | |
| stats = {"step": step} | |
| timer = profiler.Timer().tic() | |
| # Update LR from cosine schedulers. | |
| # param_groups[0] = student, param_groups[1] = aux | |
| lr_s = self.scheduler_s.get_lr() | |
| lr_a = self.scheduler_a.get_lr() | |
| stats["lr_student"] = lr_s | |
| stats["lr_aux"] = lr_a | |
| self.optimizer.param_groups[0]["lr"] = lr_s | |
| self.optimizer.param_groups[1]["lr"] = lr_a | |
| # ---------------------------------------------------------------- | |
| # Stage 1: Get tokenised batch (CPU → GPU) | |
| # ---------------------------------------------------------------- | |
| txt_ids = next(self._inf_loader) # [B, L] CPU tensor | |
| txt_ids = txt_ids.to(device, non_blocking=True) | |
| B = txt_ids.size(0) | |
| txt_uncond = None | |
| if cfg.enable_teacher_cfg: | |
| txt_uncond = self.txt_uncond_base_cpu.expand(B, -1).to(device) | |
| # # ---------------------------------------------------------------- | |
| # # Stage 2: Sample x_init ~ Uniform(K) with optional p_init mixing | |
| # # ---------------------------------------------------------------- | |
| # x_init = torch.randint(0, K, (B, T, H, W), device=device, dtype=torch.long) | |
| # if self.x_hat_prev is not None and float(cfg.p_init_mix_ratio) > 0: | |
| # n_mix = max(1, int(B * float(cfg.p_init_mix_ratio))) | |
| # x_init[:n_mix] = self.pipe.corrupt_tokens( | |
| # self.x_hat_prev[:n_mix], r=float(cfg.p_mix_corrupt_frac) | |
| # ) | |
| # ---------------------------------------------------------------- | |
| # Stage 2: Sample x_init ~ Uniform(K) with optional p_init mixing | |
| # ---------------------------------------------------------------- | |
| x_init = torch.randint(0, K, (B, T, H, W), device=device, dtype=torch.long) | |
| # 修复：使用概率触发，确保小 Batch 时模型依然能充分学习处理纯噪声 | |
| if self.x_hat_prev is not None and float(cfg.p_init_mix_ratio) > 0: | |
| if torch.rand(1).item() < float(cfg.p_init_mix_ratio): | |
| # 如果触发，只混合 batch 里的第一个样本 | |
| x_init[0] = self.pipe.corrupt_tokens( | |
| self.x_hat_prev[0:1], r=float(cfg.p_mix_corrupt_frac) | |
| ).squeeze(0) | |
| # ---------------------------------------------------------------- | |
| # Stage 3: Student 1-step on x_init — no_grad (only sample x_hat) | |
| # | |
| # Gradient-enabled forward on x_init is deferred to Stage 9b so | |
| # the KD computation graph (Stage 7, x_t) can be freed first. | |
| # ---------------------------------------------------------------- | |
| with torch.no_grad(): | |
| ids_init, rpos_init, _ = self.pipe.build_inputs( | |
| txt_ids, x_init, self.latents_shape | |
| ) | |
| logits_s_init = _get_logits( | |
| self.model("student", ids_init, rope_pos=rpos_init) | |
| ) | |
| z_s = self.pipe.extract_logits(logits_s_init, N) # [B, N, K] | |
| p_s = F.softmax(z_s / float(cfg.tau), dim=-1) # [B, N, K] | |
| x_hat = torch.multinomial(p_s.view(-1, K), 1).view(B, N) # [B, N] | |
| # if step == 1: | |
| # # 只抽 8 个 token 做 sum=1 检查，别全量 | |
| # idx = torch.randint(0, N, (8,), device=device) | |
| # p_err = (p_s[:, idx].sum(-1) - 1).abs().max().item() | |
| # assert p_err < 1e-3, f"p_s subset not normalised: {p_err}" | |
| del p_s, z_s, logits_s_init | |
| x_hat_4d = x_hat.view(B, T, H, W) | |
| # ---------------------------------------------------------------- | |
| # Stage 4: Pseudo-intermediate x_t = add_noise(x_hat, t) | |
| # ---------------------------------------------------------------- | |
| t = self.pipe.sample_t_curriculum( | |
| B, device, step, int(cfg.t_curriculum_steps) | |
| ) # [B] float ∈ (0.05, 0.995) | |
| with torch.no_grad(): | |
| x_t = self.pipe.scheduler.add_noise(x_hat_4d, t) # [B,T,H,W] long | |
| # # ---------------------------------------------------------------- | |
| # # Stage 5: Teacher forward — single [2B] forward when CFG enabled | |
| # # ---------------------------------------------------------------- | |
| # with torch.no_grad(): | |
| # if cfg.enable_teacher_cfg: | |
| # txt_dual = torch.cat([txt_ids, txt_uncond], dim=0) # [2B, L] | |
| # x_t_dual = torch.cat([x_t, x_t], dim=0) # [2B,T,H,W] | |
| # ids_dual, rpos_dual, _ = self.pipe.build_inputs( | |
| # txt_dual, x_t_dual, self.latents_shape | |
| # ) | |
| # logits_T_dual = _get_logits( | |
| # self.pipe.teacher(ids_dual, rope_pos=rpos_dual) | |
| # ) | |
| # z_T_dual = self.pipe.extract_logits(logits_T_dual, N) # [2B,N,K] | |
| # z_T_cond, z_T_uncond = z_T_dual.chunk(2, dim=0) # [B,N,K] | |
| # del logits_T_dual, z_T_dual | |
| # torch.cuda.empty_cache() | |
| # ids_t, rpos_t = ids_dual[:B], rpos_dual[:B] | |
| # else: | |
| # ids_t, rpos_t, _ = self.pipe.build_inputs( | |
| # txt_ids, x_t, self.latents_shape | |
| # ) | |
| # logits_T = _get_logits( | |
| # self.pipe.teacher(ids_t, rope_pos=rpos_t) | |
| # ) | |
| # z_T_cond = self.pipe.extract_logits(logits_T, N) # [B,N,K] | |
| # z_T_uncond = None | |
| # ids_dual, rpos_dual = ids_t, rpos_t | |
| # # CFG guided target with per-sample Bernoulli warmup. | |
| # z_T_guided = None | |
| # use_guided_ratio = 0.0 | |
| # if cfg.enable_teacher_cfg: | |
| # p_guided = _cfg_warmup_prob( | |
| # step, | |
| # float(cfg.teacher_cfg_prob), | |
| # int(cfg.teacher_cfg_warmup_steps), | |
| # ) | |
| # use_guided = torch.rand(B, device=device) < p_guided # [B] bool | |
| # use_guided_ratio = float(use_guided.float().mean().item()) | |
| # z_T_guided = _build_guided_logits( | |
| # z_T_cond, z_T_uncond, | |
| # t, float(cfg.teacher_cfg_scale), float(cfg.teacher_cfg_trunc), | |
| # ) | |
| # mask = use_guided.view(-1, 1, 1).expand_as(z_T_cond) | |
| # z_T_target = torch.where(mask, z_T_guided, z_T_cond.float()) | |
| # else: | |
| # z_T_target = z_T_cond | |
| # z_T_target = z_T_target.detach() # NO grad path to teacher | |
| # # # ---------------------------------------------------------------- | |
| # # # Stage 6: Aux update — fake_rounds iterations | |
| # # # | |
| # # # Freeze student so only aux gets gradients. With a single | |
| # # # DeepSpeed-wrapped optimizer this is the cleanest way to ensure | |
| # # # only aux params are updated. | |
| # # # ---------------------------------------------------------------- | |
| # # raw_twin = self.accelerator.unwrap_model(self.model) | |
| # # raw_twin.student.requires_grad_(False) | |
| # # raw_twin.aux.requires_grad_(True) | |
| # # loss_aux_cond_last = torch.tensor(0.0, device=device) | |
| # # loss_aux_uncond_last = torch.tensor(0.0, device=device) | |
| # # loss_aux_cond_sample_last = None | |
| # # for _fr in range(int(cfg.fake_rounds)): | |
| # # self.optimizer.zero_grad(set_to_none=True) | |
| # # if cfg.enable_teacher_cfg: | |
| # # logits_A_dual = _get_logits( | |
| # # self.model("aux", ids_dual.detach(), rope_pos=rpos_dual.detach()) | |
| # # ) | |
| # # z_A_dual = self.pipe.extract_logits(logits_A_dual, N) # [2B,N,K] | |
| # # z_A_cond, z_A_uncond = z_A_dual.chunk(2, dim=0) | |
| # # loss_aux_cond_sample = _stable_jeffrey( | |
| # # z_T_target, z_A_cond, float(cfg.tau_kd),chunk_size=1024 | |
| # # ) # [B] | |
| # # loss_aux_cond_v = loss_aux_cond_sample.mean() | |
| # # loss_aux_uncond_v = _stable_jeffrey( | |
| # # z_T_uncond.float().detach(), z_A_uncond, float(cfg.tau_kd),chunk_size=1024 | |
| # # ).mean() | |
| # # loss_aux_v = ( | |
| # # loss_aux_cond_v | |
| # # + float(cfg.lambda_kd_uncond) * loss_aux_uncond_v | |
| # # ) | |
| # # else: | |
| # # logits_A = _get_logits( | |
| # # self.model("aux", ids_t.detach(), rope_pos=rpos_t.detach()) | |
| # # ) | |
| # # z_A_cond = self.pipe.extract_logits(logits_A, N) | |
| # # loss_aux_cond_sample = _stable_jeffrey( | |
| # # z_T_target, z_A_cond, float(cfg.tau_kd),chunk_size=1024 | |
| # # ) | |
| # # loss_aux_cond_v = loss_aux_cond_sample.mean() | |
| # # loss_aux_uncond_v = torch.tensor(0.0, device=device) | |
| # # loss_aux_v = loss_aux_cond_v | |
| # # self.accelerator.backward(loss_aux_v) | |
| # # if float(cfg.grad_clip) > 0: | |
| # # torch.nn.utils.clip_grad_norm_( | |
| # # raw_twin.aux.parameters(), float(cfg.grad_clip) | |
| # # ) | |
| # # self.optimizer.step() | |
| # # self.optimizer.zero_grad(set_to_none=True) | |
| # # loss_aux_cond_last = loss_aux_cond_v.detach() | |
| # # loss_aux_uncond_last = loss_aux_uncond_v.detach() | |
| # # loss_aux_cond_sample_last = loss_aux_cond_sample.detach() # [B] | |
| # # # ---------------------------------------------------------------- | |
| # # # Stage 7: Student KD forward on x_t (with grad) | |
| # # # | |
| # # # Switch: freeze aux, unfreeze student for Stages 7-9. | |
| # # # ---------------------------------------------------------------- | |
| # # raw_twin.student.requires_grad_(True) | |
| # # raw_twin.aux.requires_grad_(False) | |
| # # self.optimizer.zero_grad(set_to_none=True) | |
| # # if cfg.enable_teacher_cfg: | |
| # # logits_S_dual = _get_logits( | |
| # # self.model("student", ids_dual.detach(), rope_pos=rpos_dual.detach()) | |
| # # ) | |
| # # z_S_dual = self.pipe.extract_logits(logits_S_dual, N) | |
| # # z_S_cond, z_S_uncond = z_S_dual.chunk(2, dim=0) | |
| # # loss_kd_cond = _stable_kl( | |
| # # z_T_target, z_S_cond, float(cfg.tau_kd), chunk_size=2048 | |
| # # ).mean() | |
| # # loss_kd_uncond = _stable_kl( | |
| # # z_T_uncond.float().detach(), z_S_uncond, float(cfg.tau_kd), chunk_size=2048 | |
| # # ).mean() | |
| # # loss_kd = loss_kd_cond + float(cfg.lambda_kd_uncond) * loss_kd_uncond | |
| # # else: | |
| # # logits_S = _get_logits( | |
| # # self.model("student", ids_t.detach(), rope_pos=rpos_t.detach()) | |
| # # ) | |
| # # z_S_cond = self.pipe.extract_logits(logits_S, N) | |
| # # loss_kd_cond = _stable_kl( | |
| # # z_T_target, z_S_cond, float(cfg.tau_kd), chunk_size=2048 | |
| # # ).mean() | |
| # # loss_kd_uncond = torch.tensor(0.0, device=device) | |
| # # loss_kd = loss_kd_cond | |
| # # # ---------------------------------------------------------------- | |
| # # # Stage 8: Reward + advantage (fully detached — no student grad) | |
| # # # | |
| # # # INVARIANT: reward and adv must never carry student gradients. | |
| # # # ---------------------------------------------------------------- | |
| # # if cfg.enable_teacher_cfg and cfg.reward_use_guided: | |
| # # z_T_for_rew = z_T_target # already detached | |
| # # else: | |
| # # z_T_for_rew = z_T_cond.detach() | |
| # # # reward[b] = -KL(z_T_cond || z_S_cond) with BOTH inputs detached | |
| # # with torch.no_grad(): | |
| # # reward = -_stable_kl( | |
| # # z_T_for_rew.detach(), z_S_cond.detach(), float(cfg.tau), chunk_size=1024 | |
| # # ) # [B] | |
| # # assert not reward.requires_grad, ( | |
| # # "[BUG] reward.requires_grad=True — student grad leaked into reward. " | |
| # # "z_S_cond must be detached before KL for reward." | |
| # # ) | |
| # # self.baseline_ema = ( | |
| # # 0.99 * self.baseline_ema + 0.01 * float(reward.mean().item()) | |
| # # ) | |
| # # adv = (reward - self.baseline_ema).detach() # [B] | |
| # # assert not adv.requires_grad, "[BUG] adv.requires_grad=True" | |
| # # # ---------------------------------------------------------------- | |
| # # # Stage 9: Two-backward student update | |
| # # # | |
| # # # 9a) KD backward first — frees the KD graph to save memory. | |
| # # # Uses no_sync() (no DDP all-reduce) so gradients are not | |
| # # # double-reduced when the PG backward syncs in 9b. | |
| # # # 9b) Fresh forward on x_init WITH grad → PG + entropy backward. | |
| # # # DDP all-reduce happens here (normal backward). | |
| # # # ---------------------------------------------------------------- | |
| # # # 9a: KD backward (no sync — first of two backwards) | |
| # # _no_sync_backward( | |
| # # self.accelerator, self.model, float(cfg.lambda_kd) * loss_kd | |
| # # ) | |
| # # # 9b: Policy + entropy — fresh forward on x_init WITH grad | |
| # # ids_init, rpos_init, _ = self.pipe.build_inputs( | |
| # # txt_ids, x_init, self.latents_shape | |
| # # ) | |
| # # logits_s_pol = _get_logits( | |
| # # self.model("student", ids_init, rope_pos=rpos_init) | |
| # # ) | |
| # # z_s_pol = self.pipe.extract_logits(logits_s_pol, N) # [B, N, K] | |
| # # logp_tok = F.log_softmax(z_s_pol / float(cfg.tau), dim=-1) # [B, N, K] | |
| # # p_s_pol = logp_tok.exp() | |
| # # # per-token average log-prob (recommended over log-prob sum) | |
| # # logp_sum = ( | |
| # # logp_tok.gather(-1, x_hat.unsqueeze(-1)).squeeze(-1).sum(-1) | |
| # # ) # [B] | |
| # # logp = logp_sum / N # [B] per-token logp | |
| # # H_mean = -(p_s_pol * logp_tok).sum(-1).mean() | |
| # # loss_pg = -(adv * logp).mean() | |
| # # lambda_ent_eff = float(cfg.lambda_ent) * (1.0 + 2.0 * use_guided_ratio) | |
| # # # Second backward: DDP all-reduce happens here. | |
| # # self.accelerator.backward( | |
| # # float(cfg.lambda_pg) * loss_pg - lambda_ent_eff * H_mean | |
| # # ) | |
| # # if float(cfg.grad_clip) > 0: | |
| # # torch.nn.utils.clip_grad_norm_( | |
| # # raw_twin.student.parameters(), float(cfg.grad_clip) | |
| # # ) | |
| # # self.optimizer.step() | |
| # # # Restore both sub-modules to trainable for next step. | |
| # # raw_twin.student.requires_grad_(True) | |
| # # raw_twin.aux.requires_grad_(True) | |
| # # # p_init mixing: store x_hat_4d (detached) for next step. | |
| # # self.x_hat_prev = x_hat_4d.detach() | |
| # # ---------------------------------------------------------------- | |
| # # Stage 6: Aux update — Fit sampled pseudo-target (x_hat) from student | |
| # # ---------------------------------------------------------------- | |
| # raw_twin = self.accelerator.unwrap_model(self.model) | |
| # raw_twin.student.requires_grad_(False) | |
| # raw_twin.aux.requires_grad_(True) | |
| # target_tokens = x_hat.detach() # [B, N] - 学生在 Stage 3 盲猜出来的画面 | |
| # for _fr in range(int(cfg.fake_rounds)): | |
| # self.optimizer.zero_grad(set_to_none=True) | |
| # if cfg.enable_teacher_cfg: | |
| # logits_A_dual = _get_logits( | |
| # self.model("aux", ids_dual.detach(), rope_pos=rpos_dual.detach()) | |
| # ) | |
| # z_A_dual = self.pipe.extract_logits(logits_A_dual, N) # [2B,N,K] | |
| # z_A_cond, z_A_uncond = z_A_dual.chunk(2, dim=0) | |
| # # Aux 拟合学生的假 token (Cross Entropy) | |
| # loss_aux_cond_v = F.cross_entropy( | |
| # z_A_cond.reshape(B * N, K), | |
| # target_tokens.reshape(B * N), | |
| # reduction="mean", | |
| # ) | |
| # loss_aux_uncond_v = F.cross_entropy( | |
| # z_A_uncond.reshape(B * N, K), | |
| # target_tokens.reshape(B * N), | |
| # reduction="mean", | |
| # ) | |
| # loss_aux_v = loss_aux_cond_v + float(cfg.lambda_kd_uncond) * loss_aux_uncond_v | |
| # else: | |
| # logits_A = _get_logits( | |
| # self.model("aux", ids_t.detach(), rope_pos=rpos_t.detach()) | |
| # ) | |
| # z_A_cond = self.pipe.extract_logits(logits_A, N) | |
| # loss_aux_cond_v = F.cross_entropy( | |
| # z_A_cond.reshape(B * N, K), | |
| # target_tokens.reshape(B * N), | |
| # reduction="mean", | |
| # ) | |
| # loss_aux_uncond_v = torch.tensor(0.0, device=device) | |
| # loss_aux_v = loss_aux_cond_v | |
| # self.accelerator.backward(loss_aux_v) | |
| # if float(cfg.grad_clip) > 0: | |
| # torch.nn.utils.clip_grad_norm_( | |
| # raw_twin.aux.parameters(), float(cfg.grad_clip) | |
| # ) | |
| # self.optimizer.step() | |
| # loss_aux_cond_last = loss_aux_cond_v.detach() | |
| # # ---------------------------------------------------------------- | |
| # # Stage 7 & 8: Student KD update & Aux Bridge (Gradient Injection) | |
| # # ---------------------------------------------------------------- | |
| # raw_twin.student.requires_grad_(True) | |
| # raw_twin.aux.requires_grad_(False) | |
| # self.optimizer.zero_grad(set_to_none=True) | |
| # # 7a. Student KD forward on x_t (保持原样) | |
| # if cfg.enable_teacher_cfg: | |
| # logits_S_dual = _get_logits( | |
| # self.model("student", ids_dual.detach(), rope_pos=rpos_dual.detach()) | |
| # ) | |
| # z_S_dual = self.pipe.extract_logits(logits_S_dual, N) | |
| # z_S_cond, z_S_uncond = z_S_dual.chunk(2, dim=0) | |
| # # --- [新增] 立刻释放显存 --- | |
| # del logits_S_dual, z_S_dual | |
| # loss_kd_cond = _stable_kl( | |
| # z_T_target, z_S_cond, float(cfg.tau_kd), chunk_size=256 #2048 | |
| # ).mean() | |
| # loss_kd_uncond = _stable_kl( | |
| # z_T_uncond.float().detach(), z_S_uncond, float(cfg.tau_kd), chunk_size=256 #2048 | |
| # ).mean() | |
| # loss_kd = loss_kd_cond + float(cfg.lambda_kd_uncond) * loss_kd_uncond | |
| # else: | |
| # logits_S = _get_logits( | |
| # self.model("student", ids_t.detach(), rope_pos=rpos_t.detach()) | |
| # ) | |
| # z_S_cond = self.pipe.extract_logits(logits_S, N) | |
| # loss_kd_cond = _stable_kl( | |
| # z_T_target, z_S_cond, float(cfg.tau_kd), chunk_size=256 #2048 | |
| # ).mean() | |
| # loss_kd_uncond = torch.tensor(0.0, device=device) | |
| # loss_kd = loss_kd_cond | |
| # # 7b. 获取 Aux 的预测 (无梯度) 作为计算桥梁 | |
| # with torch.no_grad(): | |
| # if cfg.enable_teacher_cfg: | |
| # logits_A_dual = _get_logits( | |
| # self.model("aux", ids_dual.detach(), rope_pos=rpos_dual.detach()) | |
| # ) | |
| # z_A_dual = self.pipe.extract_logits(logits_A_dual, N) | |
| # z_A_cond, _ = z_A_dual.chunk(2, dim=0) | |
| # # --- [新增] 立刻释放显存 --- | |
| # del logits_A_dual, z_A_dual | |
| # else: | |
| # logits_A = _get_logits( | |
| # self.model("aux", ids_t.detach(), rpos_t.detach()) | |
| # ) | |
| # z_A_cond = self.pipe.extract_logits(logits_A, N) | |
| # # 8. Student 对初始噪声 x_init 进行带梯度的前向传播 | |
| # ids_init, rpos_init, _ = self.pipe.build_inputs( | |
| # txt_ids, x_init, self.latents_shape | |
| # ) | |
| # logits_s_pol = _get_logits( | |
| # self.model("student", ids_init, rope_pos=rpos_init) | |
| # ) | |
| # z_s_pol = self.pipe.extract_logits(logits_s_pol, N) | |
| # # --- 核心数学修正：将 Logits 转换为概率，防止梯度爆炸 --- | |
| # p_T = F.softmax(z_T_target / float(cfg.tau_kd), dim=-1) | |
| # p_A = F.softmax(z_A_cond / float(cfg.tau_kd), dim=-1) | |
| # # 目标方向：Teacher 概率 - Aux 概率 (遵循论文公式推导) | |
| # bridge_target = (p_T - p_A).detach() | |
| # # 利用 MSE Trick 强制注入梯度 | |
| # loss_bridge = 0.5 * F.mse_loss( | |
| # z_s_pol.float(), | |
| # (z_s_pol.float() + bridge_target).detach() | |
| # ) | |
| # # 9. 单次反向传播 (合并 KD 和 Bridge) | |
| # # 借用原来的 lambda_pg 参数来控制 bridge 损失的权重 | |
| # loss_student = float(cfg.lambda_kd) * loss_kd + float(cfg.lambda_pg) * loss_bridge | |
| # self.accelerator.backward(loss_student) | |
| # if float(cfg.grad_clip) > 0: | |
| # torch.nn.utils.clip_grad_norm_( | |
| # raw_twin.student.parameters(), float(cfg.grad_clip) | |
| # ) | |
| # self.optimizer.step() | |
| # # 恢复两者的可训练状态 | |
| # raw_twin.student.requires_grad_(True) | |
| # raw_twin.aux.requires_grad_(True) | |
| # # --- 兼容原始日志输出的占位符 --- | |
| # H_mean = torch.tensor(0.0, device=device) | |
| # loss_pg = loss_bridge.detach() # 将 bridge 损失映射给 pg 显示 | |
| # logp = torch.tensor(0.0, device=device) | |
| # self.baseline_ema = 0.0 | |
| # ---------------------------------------------------------------- | |
| # Stage 5: Teacher forward — 破除视图死锁，生成目标后立刻释放 | |
| # ---------------------------------------------------------------- | |
| with torch.no_grad(): | |
| if cfg.enable_teacher_cfg: | |
| txt_dual = torch.cat([txt_ids, txt_uncond], dim=0) # [2B, L] | |
| x_t_dual = torch.cat([x_t, x_t], dim=0) # [2B,T,H,W] | |
| ids_dual, rpos_dual, _ = self.pipe.build_inputs( | |
| txt_dual, x_t_dual, self.latents_shape | |
| ) | |
| logits_T_dual = _get_logits( | |
| self.pipe.teacher(ids_dual, rope_pos=rpos_dual) | |
| ) | |
| z_T_dual = self.pipe.extract_logits(logits_T_dual, N) # [2B,N,K] | |
| # 【显存救星 1】使用 .clone() 打断视图依赖，使得原始巨型张量可以被回收 | |
| z_T_cond = z_T_dual[0:1].clone() # [1,N,K] | |
| z_T_uncond = z_T_dual[1:2].clone() # [1,N,K] | |
| ids_t, rpos_t = ids_dual[:B], rpos_dual[:B] | |
| # 立刻释放 17 GB 的双路缓冲 | |
| del logits_T_dual, z_T_dual | |
| torch.cuda.empty_cache() | |
| else: | |
| ids_t, rpos_t, _ = self.pipe.build_inputs(txt_ids, x_t, self.latents_shape) | |
| logits_T = _get_logits(self.pipe.teacher(ids_t, rope_pos=rpos_t)) | |
| z_T_cond = self.pipe.extract_logits(logits_T, N) | |
| z_T_uncond = None | |
| # 计算 CFG guided target | |
| z_T_guided = None | |
| use_guided_ratio = 0.0 | |
| if cfg.enable_teacher_cfg: | |
| p_guided = _cfg_warmup_prob(step, float(cfg.teacher_cfg_prob), int(cfg.teacher_cfg_warmup_steps)) | |
| use_guided = torch.rand(B, device=device) < p_guided | |
| use_guided_ratio = float(use_guided.float().mean().item()) | |
| z_T_guided = _build_guided_logits( | |
| z_T_cond, z_T_uncond, | |
| t, float(cfg.teacher_cfg_scale), float(cfg.teacher_cfg_trunc), | |
| ) | |
| mask = use_guided.view(-1, 1, 1).expand_as(z_T_cond) | |
| # 【显存救星 2】保持为 bf16 类型，避免膨胀到 8.5GB | |
| z_T_target = torch.where(mask, z_T_guided, z_T_cond).to(dtype=z_T_cond.dtype).detach() | |
| # 立刻清理所有中间推导变量 | |
| del z_T_cond, z_T_uncond, z_T_guided | |
| torch.cuda.empty_cache() | |
| else: | |
| z_T_target = z_T_cond.detach() | |
| # ---------------------------------------------------------------- | |
| # Stage 6: Aux update — 【显存救星 3】强行降维为单路前向传播 (Batch=1) | |
| # ---------------------------------------------------------------- | |
| raw_twin = self.accelerator.unwrap_model(self.model) | |
| raw_twin.student.requires_grad_(False) | |
| raw_twin.aux.requires_grad_(True) | |
| target_tokens = x_hat.detach() | |
| for _fr in range(int(cfg.fake_rounds)): | |
| self.optimizer.zero_grad(set_to_none=True) | |
| # 只处理单路 ids_t，不处理 dual，砍掉 Aux 50% 显存！ | |
| logits_A = _get_logits( | |
| self.model("aux", ids_t.detach(), rope_pos=rpos_t.detach()) | |
| ) | |
| z_A_cond = self.pipe.extract_logits(logits_A, N) | |
| loss_aux_cond_v = F.cross_entropy( | |
| z_A_cond.reshape(B * N, K), | |
| target_tokens.reshape(B * N), | |
| reduction="mean", | |
| ) | |
| self.accelerator.backward(loss_aux_cond_v) | |
| if float(cfg.grad_clip) > 0: | |
| torch.nn.utils.clip_grad_norm_(raw_twin.aux.parameters(), float(cfg.grad_clip)) | |
| self.optimizer.step() | |
| # 必须立刻释放 | |
| del logits_A, z_A_cond | |
| torch.cuda.empty_cache() | |
| loss_aux_cond_last = loss_aux_cond_v.detach() | |
| # ---------------------------------------------------------------- | |
| # Stage 7 & 8: Student KD update & Aux Bridge | |
| # ---------------------------------------------------------------- | |
| raw_twin.student.requires_grad_(True) | |
| raw_twin.aux.requires_grad_(False) | |
| self.optimizer.zero_grad(set_to_none=True) | |
| # 7a. Student KD (强行降维为单路前向传播 Batch=1) | |
| logits_S = _get_logits( | |
| self.model("student", ids_t.detach(), rope_pos=rpos_t.detach()) | |
| ) | |
| z_S_cond = self.pipe.extract_logits(logits_S, N) | |
| # 使用 128 chunk size，确保极致安全 | |
| loss_kd = _stable_kl( | |
| z_T_target, z_S_cond, float(cfg.tau_kd), chunk_size=128 | |
| ).mean() | |
| del logits_S, z_S_cond | |
| torch.cuda.empty_cache() | |
| # 7b. 获取 Aux 的预测作为桥梁 (依然单路) | |
| with torch.no_grad(): | |
| logits_A = _get_logits( | |
| self.model("aux", ids_t.detach(), rope_pos=rpos_t.detach()) | |
| ) | |
| z_A_cond = self.pipe.extract_logits(logits_A, N) | |
| # 8. Student 对 x_init 进行前向传播 | |
| ids_init, rpos_init, _ = self.pipe.build_inputs(txt_ids, x_init, self.latents_shape) | |
| logits_s_pol = _get_logits( | |
| self.model("student", ids_init, rope_pos=rpos_init) | |
| ) | |
| z_s_pol = self.pipe.extract_logits(logits_s_pol, N) | |
| # 【显存救星 4】在 bf16 精度下计算 Softmax 概率，防止 float32 炸存 | |
| p_T = F.softmax(z_T_target / float(cfg.tau_kd), dim=-1).to(z_s_pol.dtype) | |
| p_A = F.softmax(z_A_cond / float(cfg.tau_kd), dim=-1).to(z_s_pol.dtype) | |
| bridge_target = (p_T - p_A).detach() | |
| # 拿到 bridge_target 后，前面所有百兆甚至 G 级的张量统统干掉 | |
| del p_T, p_A, logits_A, z_A_cond, z_T_target | |
| torch.cuda.empty_cache() | |
| # 伪梯度注入 | |
| loss_bridge = 0.5* K * F.mse_loss( | |
| z_s_pol.float(), | |
| (z_s_pol.float() + bridge_target.float()).detach() | |
| ) | |
| # 9. 统一反向传播 | |
| loss_student = float(cfg.lambda_kd) * loss_kd + float(cfg.lambda_pg) * loss_bridge | |
| self.accelerator.backward(loss_student) | |
| if float(cfg.grad_clip) > 0: | |
| torch.nn.utils.clip_grad_norm_(raw_twin.student.parameters(), float(cfg.grad_clip)) | |
| self.optimizer.step() | |
| # 恢复状态 | |
| raw_twin.student.requires_grad_(True) | |
| raw_twin.aux.requires_grad_(True) | |
| # 最后的清理 | |
| del logits_s_pol, z_s_pol, bridge_target | |
| torch.cuda.empty_cache() | |
| H_mean = torch.tensor(0.0, device=device) | |
| loss_pg = loss_bridge.detach() | |
| logp = torch.tensor(0.0, device=device) | |
| self.baseline_ema = 0.0 | |
| # Advance LR schedulers. | |
| self.scheduler_s.step() | |
| self.scheduler_a.step() | |
| # ---------------------------------------------------------------- | |
| # Step 1 sanity assertions (lightweight; runs only at step 1) | |
| # ---------------------------------------------------------------- | |
| # if step == 1: | |
| # self._step1_assertions( | |
| # x_init, ids_init, rpos_init, z_s, p_s, logp, | |
| # z_T_cond, z_S_cond, x_t, B, T, H, W, | |
| # ) | |
| # ---------------------------------------------------------------- | |
| # Token-level collapse detection | |
| # ---------------------------------------------------------------- | |
| tok_entropy = self._token_entropy(x_hat) | |
| if not hasattr(self, "_init_tok_entropy"): | |
| self._init_tok_entropy = tok_entropy | |
| collapse_frac = float(cfg.get("collapse_warn_frac", 0.2)) | |
| if tok_entropy < collapse_frac * self._init_tok_entropy: | |
| self.logger.warning( | |
| f"[COLLAPSE] step={step} tok_H={tok_entropy:.3f} " | |
| f"init={self._init_tok_entropy:.3f} " | |
| f"ratio={tok_entropy / max(self._init_tok_entropy, 1e-8):.2f} " | |
| f"< {collapse_frac}. Try increasing lambda_ent." | |
| ) | |
| stats["time"] = timer.toc() | |
| stats["metrics"] = collections.OrderedDict( | |
| sorted( | |
| { | |
| "loss_aux_cond": float(loss_aux_cond_last.item()), | |
| "loss_kd_cond": float(loss_kd.item()), | |
| "loss_pg": float(loss_pg.item()), | |
| "H_mean": float(H_mean.item()), | |
| "tok_entropy": float(tok_entropy), | |
| "mean_logp_tok": float(logp.mean().item()), | |
| "baseline_ema": float(self.baseline_ema), | |
| "use_guided_ratio": float(use_guided_ratio), | |
| }.items() | |
| ) | |
| ) | |
| return stats | |
| # ----------------------------------------------------------------------- | |
| # Train loop | |
| # ----------------------------------------------------------------------- | |
| def train_loop(self): | |
| """Main training loop (mirrors diffnext.engine.train_engine.Trainer).""" | |
| cfg_exp = self.config.experiment | |
| max_steps = int(self.config.training.max_train_steps) | |
| log_every = int(cfg_exp.log_every) | |
| save_every = int(cfg_exp.save_every) | |
| self.global_step = int(self.config.experiment.get("resume_iter", 0)) | |
| # Sync LR schedulers to resume step (set _step_count directly; | |
| # CosineLR uses _step_count internally in get_decay()). | |
| self.scheduler_s._step_count = self.global_step | |
| self.scheduler_a._step_count = self.global_step | |
| # [可选补充] 如果是续传，让 accelerator 自动恢复被切分的 Optimizer 等状态 | |
| if self.global_step > 0: | |
| ckpt_dir = os.path.join(self.config.experiment.output_dir, "checkpoints", f"checkpoint-{self.global_step}") | |
| if os.path.exists(ckpt_dir): | |
| self.accelerator.load_state(ckpt_dir) | |
| self.logger.info(f"✅ ZeRO-3 完整状态 (包含 Optimizer) 已从 {ckpt_dir} 恢复") | |
| timer = profiler.Timer() | |
| self.logger.info( | |
| f"[train] Starting from step {self.global_step} / {max_steps}" | |
| ) | |
| while self.global_step < max_steps: | |
| self.global_step += 1 | |
| with timer.tic_and_toc(): | |
| stats = self.run_step(self.global_step) | |
| self._add_metrics(stats) | |
| if self.global_step % log_every == 0: | |
| self._log_metrics(stats) | |
| if self.global_step % (10 * log_every) == 0: | |
| self.logger.info( | |
| profiler.get_progress(timer, self.global_step, max_steps) | |
| ) | |
| if self.global_step % save_every == 0: | |
| self.save(self.global_step) | |
| # Final log + save (only when loop ran at least one step). | |
| if self.global_step > int(self.config.experiment.get("resume_iter", 0)): | |
| self._log_metrics({**stats, "step": self.global_step}) # noqa: F821 | |
| self.accelerator.wait_for_everyone() | |
| self.save(self.global_step, suffix="final") | |
| self.accelerator.end_training() | |
| # ----------------------------------------------------------------------- | |
| # Checkpoint helpers | |
| # ----------------------------------------------------------------------- | |
| # def save(self, step: int, suffix: str = None) -> None: | |
| # """Save student + aux state_dicts (rank0 only). | |
| # Saved as: | |
| # <output_dir>/checkpoints/checkpoint-<step>/student.pt | |
| # <output_dir>/checkpoints/checkpoint-<step>/aux.pt | |
| # The student.pt can be used for inference by replacing the | |
| # transformer weights in a URSAPipeline (see README). | |
| # """ | |
| # if not self.accelerator.is_main_process: | |
| # return | |
| # folder = f"checkpoint-{suffix}" if suffix else f"checkpoint-{step}" | |
| # ckpt_dir = os.path.join( | |
| # self.config.experiment.output_dir, "checkpoints", folder | |
| # ) | |
| # os.makedirs(ckpt_dir, exist_ok=True) | |
| # raw_student = self.accelerator.unwrap_model(self.model).student | |
| # raw_aux = self.accelerator.unwrap_model(self.model).aux | |
| # student_path = os.path.join(ckpt_dir, "student.pt") | |
| # aux_path = os.path.join(ckpt_dir, "aux.pt") | |
| # torch.save(raw_student.state_dict(), student_path) | |
| # torch.save(raw_aux.state_dict(), aux_path) | |
| # # Also save training state for resuming. | |
| # state = { | |
| # "global_step": step, | |
| # "baseline_ema": self.baseline_ema, | |
| # "optimizer": self.optimizer.state_dict(), | |
| # } | |
| # torch.save(state, os.path.join(ckpt_dir, "train_state.pt")) | |
| # self.logger.info(f"[save] step={step} → {ckpt_dir}") | |
| def save(self, step: int, suffix: str = None) -> None: | |
| """Save student + aux state_dicts (支持 DeepSpeed ZeRO-3 自动聚合).""" | |
| # ⚠️ 【极其重要】：get_state_dict 必须由所有 8 张卡共同执行！ | |
| # 绝对不能把它放在 is_main_process 判断的里面，否则会触发跨卡死锁！ | |
| full_state_dict = self.accelerator.get_state_dict(self.model) | |
| # 只有主进程（0号卡）负责把聚合好的完整参数写进硬盘 | |
| if not self.accelerator.is_main_process: | |
| return | |
| folder = f"checkpoint-{suffix}" if suffix else f"checkpoint-{step}" | |
| ckpt_dir = os.path.join( | |
| self.config.experiment.output_dir, "checkpoints", folder | |
| ) | |
| os.makedirs(ckpt_dir, exist_ok=True) | |
| # 从 TwinModel 的完整字典中，根据前缀拆分出 student 和 aux 的独立权重 | |
| student_state = {k.replace("student.", ""): v for k, v in full_state_dict.items() if k.startswith("student.")} | |
| aux_state = {k.replace("aux.", ""): v for k, v in full_state_dict.items() if k.startswith("aux.")} | |
| student_path = os.path.join(ckpt_dir, "student.pt") | |
| aux_path = os.path.join(ckpt_dir, "aux.pt") | |
| torch.save(student_state, student_path) | |
| torch.save(aux_state, aux_path) | |
| # 保存辅助状态 | |
| state = { | |
| "global_step": step, | |
| "baseline_ema": self.baseline_ema, | |
| } | |
| torch.save(state, os.path.join(ckpt_dir, "train_state.pt")) | |
| self.logger.info(f"[save] step={step} → {ckpt_dir} (ZeRO-3 Gathered)") | |
| # ----------------------------------------------------------------------- | |
| # Logging helpers (same API as original Trainer) | |
| # ----------------------------------------------------------------------- | |
| def _add_metrics(self, stats: dict) -> None: | |
| for k, v in stats["metrics"].items(): | |
| if k not in self.metrics: | |
| self.metrics[k] = profiler.SmoothedValue() | |
| self.metrics[k].update(v) | |
| def _log_metrics(self, stats: dict) -> None: | |
| iter_template = "Iteration %d, lr_s=%.2e lr_a=%.2e, time=%.2fs" | |
| self.logger.info( | |
| iter_template | |
| % ( | |
| stats["step"], | |
| stats.get("lr_student", 0.0), | |
| stats.get("lr_aux", 0.0), | |
| stats.get("time", 0.0), | |
| ) | |
| ) | |
| metric_template = " Train %s: %s" | |
| for k, v in self.metrics.items(): | |
| self.logger.info(metric_template % (k, v)) | |
| tracker_logs = {k: v.median for k, v in self.metrics.items()} | |
| tracker_logs.update( | |
| { | |
| "lr_student": stats.get("lr_student", 0.0), | |
| "time": stats.get("time", 0.0), | |
| } | |
| ) | |
| self.accelerator.log(tracker_logs, step=stats["step"]) | |
| self.metrics.clear() | |
| # ----------------------------------------------------------------------- | |
| # Sanity checks (step 1 only) | |
| # ----------------------------------------------------------------------- | |
| def _step1_assertions( | |
| self, x_init, ids_init, rpos_init, z_s, p_s, logp, | |
| z_T_cond, z_S_cond, x_t, B, T, H, W, | |
| ) -> None: | |
| """Shape / value-domain assertions (mirrors single-card script).""" | |
| N, K = self.N, self.K | |
| lm_vocab = self.pipe.teacher.config.lm_vocab_size | |
| L_plus_N1 = ids_init.size(1) | |
| txt_len = L_plus_N1 - (N + 1) | |
| assert x_init.dtype == torch.long | |
| assert x_init.min() >= 0 and x_init.max() < K | |
| assert ids_init.shape == (B, L_plus_N1), ids_init.shape | |
| txt_part = ids_init[:, :txt_len] | |
| vis_part = ids_init[:, -N:] | |
| assert (txt_part < lm_vocab).all(), "text tokens in visual range" | |
| assert (vis_part >= lm_vocab).all(), "visual tokens not shifted" | |
| assert (vis_part < lm_vocab + K).all(), "visual tokens exceed lm_vocab+K" | |
| assert rpos_init.shape == (B, L_plus_N1, 3), rpos_init.shape | |
| assert z_s.shape == (B, N, K), z_s.shape | |
| p_err = float((p_s.sum(-1) - 1).abs().max().item()) | |
| assert p_err < 1e-3, f"p_s not normalised: max_dev={p_err:.2e}" | |
| assert not torch.isnan(logp).any(), "logp has NaN" | |
| assert not torch.isinf(logp).any(), "logp has Inf" | |
| assert x_t.min() >= 0 and x_t.max() < K | |
| assert z_T_cond.shape == z_S_cond.shape == (B, N, K), ( | |
| f"z_T_cond={z_T_cond.shape} z_S_cond={z_S_cond.shape}" | |
| ) | |
| # Teacher has no grad. | |
| teacher_grads = [ | |
| p for p in self.pipe.teacher.parameters() if p.grad is not None | |
| ] | |
| assert len(teacher_grads) == 0, "teacher has grads — not frozen" | |
| # Student has grad (from PG backward). | |
| raw_s = self.accelerator.unwrap_model(self.model).student | |
| student_grad_norms = [ | |
| float(p.grad.norm().item()) | |
| for p in raw_s.parameters() | |
| if p.grad is not None | |
| ] | |
| assert len(student_grad_norms) > 0, "student has NO grads — grad flow broken" | |
| # ########################## | |
| # raw_t = self.pipe.teacher | |
| # raw_s = self.accelerator.unwrap_model(self.model).student | |
| # # (a) 共享存储检查：零开销 | |
| # pt0 = next(raw_t.parameters()) | |
| # ps0 = next(raw_s.parameters()) | |
| # self.logger.info(f"[assert] shared_storage={pt0.data_ptr() == ps0.data_ptr()}") | |
| # # (b) 参数差异：只采样前 4096 个元素，避免巨型临时张量 | |
| # with torch.no_grad(): | |
| # a = pt0.view(-1)[:4096].float() | |
| # b = ps0.view(-1)[:4096].float() | |
| # self.logger.info(f"[assert] param_delta_sample_max={float((a-b).abs().max().item()):.3e}") | |
| # # (c) logits 差异：只采样小子块（64 token × 256 vocab） | |
| # with torch.no_grad(): | |
| # idx_n = torch.randint(0, self.N, (64,), device=z_T_cond.device) | |
| # idx_k = torch.randint(0, self.K, (256,), device=z_T_cond.device) | |
| # subT = z_T_cond[0, idx_n][:, idx_k].float() | |
| # subS = z_S_cond[0, idx_n][:, idx_k].float() | |
| # self.logger.info(f"[assert] logits_delta_sub_max={float((subT-subS).abs().max().item()):.3e}") | |
| # ########################### | |
| self.logger.info("[assert] Step-1 shape/grad assertions PASSED ✓") | |
| self.logger.info( | |
| f"[assert] z_T_cond shape={z_T_cond.shape} " | |
| f"min={float(z_T_cond.min().item()):.3f} " | |
| f"max={float(z_T_cond.max().item()):.3f}" | |
| ) | |
| self.logger.info( | |
| f"[assert] z_S_cond shape={z_S_cond.shape} " | |
| f"min={float(z_S_cond.min().item()):.3f} " | |
| f"max={float(z_S_cond.max().item()):.3f}" | |
| ) | |
| def _token_entropy(x_hat: torch.Tensor) -> float: | |
| """Histogram entropy of sampled token indices (collapse detection).""" | |
| counts = x_hat.flatten().bincount(minlength=1).float() | |
| p = counts / counts.sum() | |
| p = p[p > 0] | |
| return float(-(p * p.log()).sum().item()) | |
| def main(): | |
| """Entry point — identical pattern to scripts/train.py.""" | |
| config = omegaconf_utils.get_config() | |
| os.makedirs(config.experiment.output_dir, exist_ok=True) | |
| accelerator = accelerate_utils.build_accelerator(config) | |
| accelerate_utils.build_wandb(config, accelerator=accelerator) | |
| logger = accelerate_utils.set_logger( | |
| config.experiment.output_dir, accelerator=accelerator | |
| ) | |
| device_seed = int(config.training.seed) + accelerator.process_index | |
| engine_utils.manual_seed(device_seed, (accelerator.device.index, device_seed)) | |
| if accelerator.is_main_process: | |
| config_path = os.path.join(config.experiment.output_dir, "config.yaml") | |
| omegaconf_utils.save_config(config, config_path) | |
| logger.info(f"Config:\n{omegaconf_utils.config_to_yaml(config)}") | |
| trainer = DistillTrainer(config, accelerator, logger) | |
| trainer.train_loop() | |
| if __name__ == "__main__": | |
| main() | |