1 line
67 KiB
JSON
1 line
67 KiB
JSON
{"repo": "byliutao/cdm", "files": {"scripts/train.py": {"bacc57f60c4366cb9374b59860fb094e4edb419f": {"t": 1785674300.942355, "n": 66450, "h": "da4bf4b4ce7c268e", "c": "from collections import defaultdict\nimport os\nimport json\nimport datetime\nfrom concurrent import futures\nimport time\nimport math\nfrom functools import partial\nfrom absl import app, flags\nfrom PIL import Image\n\nimport torch\nfrom torch.utils.data import DataLoader\nfrom torch.utils.data.distributed import DistributedSampler\nfrom ml_collections import config_flags\nimport tqdm\nfrom accelerate import Accelerator\nfrom accelerate.utils import ProjectConfiguration, set_seed as accelerate_set_seed, DistributedType\nfrom accelerate.logging import get_logger\n\nimport cdm.rewards.rewards\nfrom cdm.models.ema import EMAModuleWrapper\nfrom cdm.utils.logger import create_logger, log_fake_teacher_metrics\nfrom cdm.models.model_adapters import create_model_adapter, wrap_longcat_transformer_for_fsdp\nfrom cdm.pipeline.pipeline import pipeline_infer\nfrom cdm.utils.common import (\n set_seed, _repeat_embeds, sample_sigmas, enable_full_determinism,\n get_mixed_precision_dtype, create_autocast_context,\n)\nfrom cdm.data.datasets import DATASET_REGISTRY, DistributedKRepeatSampler\nfrom cdm.evaluation import eval_distillation_fn, eval_fake_teacher_viz_fn, shutdown_image_save_executor\nfrom cdm.training import (\n compute_cfg_loss,\n compute_fake_teacher_diffusion_loss, compute_dm_loss,\n compute_batch_mean_var_kl_loss,\n ModelMode, ModelComponents, get_model_mode, create_frozen_model,\n LORA_TARGET_MODULES, initialize_models, is_student_lora, is_ft_lora, has_fake_teacher,\n ema_merge_ft_with_student,\n ModelManager, patch_model_for_pipeline,\n save_fsdp_full_checkpoint, compute_grad_stats, create_lr_lambda, build_fsdp_plugin_if_needed,\n)\n\ntqdm = partial(tqdm.tqdm, dynamic_ncols=True)\n\nFLAGS = flags.FLAGS\nconfig_flags.DEFINE_config_file(\"config\", \"config/config.py\", \"Training configuration.\")\nflags.DEFINE_string(\"unique_id\", \"-1\", \"Unique experiment ID passed from submit.py. '-1' means not assigned.\")\n\nlogger = get_logger(__name__)\n\n\ndef extract_traj(data_list, indices_cpu):\n \"\"\"Extract per-sample data from a trajectory list using per-sample indices.\"\"\"\n return torch.stack([data_list[idx][i] for i, idx in enumerate(indices_cpu)])\n\n\ndef make_indices(shared_cpu, batch_size, num_steps, device):\n \"\"\"Return per-sample trajectory indices as (tensor, cpu_list).\n \"\"\"\n if shared_cpu is not None:\n idx_cpu = shared_cpu\n idx_tensor = torch.tensor(idx_cpu, device=device, dtype=torch.long)\n else:\n idx_tensor = torch.randint(0, num_steps, (batch_size,), device=device)\n idx_cpu = idx_tensor.cpu().tolist()\n return idx_tensor, idx_cpu\n\n\ndef main(_):\n config = FLAGS.config\n assert config.sample.guidance_scale == 1.0\n\n mixed_precision = config.mixed_precision or \"no\"\n \n # Unified experiment directory: logs and checkpoints live together\n experiment_dir = os.path.join(config.output_dir, config.run_name)\n log_dir = os.path.join(experiment_dir, \"tensorboard\")\n checkpoint_dir = os.path.join(experiment_dir, \"checkpoints\")\n \n # Build FSDP plugin dynamically based on model type (only when using FSDP backend)\n fsdp_plugin = build_fsdp_plugin_if_needed(config)\n \n accelerator = Accelerator(\n mixed_precision=mixed_precision,\n gradient_accumulation_steps=config.train.gradient_accumulation_steps,\n project_config=ProjectConfiguration(project_dir=experiment_dir, logging_dir=log_dir),\n fsdp_plugin=fsdp_plugin,\n )\n \n device = accelerator.device\n rank, world_size = accelerator.process_index, accelerator.num_processes\n \n logger.info(f\"Mixed precision: {mixed_precision}\", main_process_only=True)\n\n # Logger Init\n exp_logger = None\n if accelerator.is_main_process:\n os.makedirs(log_dir, exist_ok=True)\n os.makedirs(checkpoint_dir, exist_ok=True)\n config_dict = config.to_dict()\n config_dict.update({\"world_size\": world_size})\n \n # Save config snapshot for reproducibility\n config_save_path = os.path.join(experiment_dir, \"config.json\")\n with open(config_save_path, \"w\") as f:\n json.dump(config_dict, f, indent=2, default=str)\n logger.info(f\"Saved config to {config_save_path}\")\n \n exp_logger = create_logger(\n log_dir=log_dir, run_name=config.run_name, config=config_dict,\n logger_type=config.logger_type,\n project=config.wandb_project,\n entity=config.wandb_entity, api_key=config.wandb_api_key,\n )\n \n # ==================== Seed & Determinism ====================\n if config.deterministic:\n # Full deterministic mode: reproducible across runs (slower due to deterministic algorithms)\n # Each rank uses seed+rank so different GPUs produce different noise, but the same\n # rank always produces the same sequence across runs.\n enable_full_determinism(config.seed + rank)\n torch.backends.cuda.matmul.allow_tf32 = False\n torch.backends.cudnn.allow_tf32 = False\n logger.info(f\"🔒 Deterministic mode ON: seed={config.seed}, rank_seed={config.seed + rank}, TF32=disabled\", main_process_only=True)\n else:\n # Performance mode: set seed for basic reproducibility, allow non-deterministic algorithms\n accelerate_set_seed(config.seed + rank, deterministic=False)\n set_seed(config.seed, rank=rank, deterministic=False)\n if config.allow_tf32:\n torch.backends.cuda.matmul.allow_tf32 = True\n torch.backends.cudnn.allow_tf32 = True\n logger.info(f\"⚡ Performance mode: seed={config.seed}, rank_seed={config.seed + rank}, TF32={config.allow_tf32}\", main_process_only=True)\n\n mixed_precision_dtype = get_mixed_precision_dtype(config.mixed_precision)\n frozen_model_dtype = mixed_precision_dtype or torch.float32\n text_encoder_dtype = mixed_precision_dtype or torch.float32\n train_config = config.train\n student_config = train_config.student\n # Fake teacher node may be absent in config (FT feature can be turned off entirely).\n fake_teacher_config = getattr(train_config, 'fake_teacher', None)\n fake_teacher_enabled = fake_teacher_config is not None and fake_teacher_config.enabled\n # if not fake_teacher_enabled:\n # raise ValueError(\"Fake teacher must be enabled. Set train.fake_teacher.enabled=True in config.\")\n student_use_lora = student_config.use_lora\n\n # ==================== Create Model Adapter ====================\n base_model_type = config.base_model\n model_adapter = create_model_adapter(base_model_type, config)\n model_type_str = model_adapter.get_model_type() # e.g. \"sd3\" or \"longcat\"\n logger.info(f\"Model adapter: {base_model_type} (pipeline model_type={model_type_str})\", main_process_only=True)\n\n # ==================== Load Pipeline ====================\n pipeline_kwargs = {}\n if config.attn_implementation == \"flash_attention_2\":\n pipeline_kwargs[\"torch_dtype\"] = mixed_precision_dtype or torch.bfloat16\n pipeline = model_adapter.load_pipeline(config.pretrained.model, **pipeline_kwargs)\n\n if hasattr(pipeline, 'safety_checker'):\n pipeline.safety_checker = None\n pipeline.set_progress_bar_config(position=1, disable=not accelerator.is_main_process, leave=False, desc=\"Timestep\", dynamic_ncols=True)\n \n # Freeze non-trainable components\n model_adapter.freeze_non_trainable(pipeline)\n pipeline.transformer.requires_grad_(not student_use_lora)\n \n text_encoders = model_adapter.get_text_encoders(pipeline)\n tokenizers = model_adapter.get_tokenizers(pipeline)\n\n # Move models to device\n pipeline.vae.to(device, dtype=torch.float32)\n for enc in text_encoders:\n enc.to(device, dtype=text_encoder_dtype)\n transformer = pipeline.transformer.to(device)\n\n # Memory optimizations\n if train_config.enable_xformers:\n try:\n pipeline.transformer.enable_xformers_memory_efficient_attention()\n pipeline.vae.enable_xformers_memory_efficient_attention()\n logger.info(\"😊 [Memory优化] xformers enabled\", main_process_only=True)\n except Exception as e:\n logger.warning(f\"Failed to enable xformers: {e}\")\n if train_config.vae_slicing:\n pipeline.vae.enable_slicing()\n logger.info(\"😊 [Memory优化] VAE slicing enabled\", main_process_only=True)\n\n # ==================== Model Initialization ====================\n # Guard: gradient_accumulation_steps > 1 is not supported with fake teacher training.\n # FT and Student share the same accelerator accumulation counter, causing sync timing conflicts.\n if train_config.gradient_accumulation_steps > 1:\n raise ValueError(\n f\"gradient_accumulation_steps={train_config.gradient_accumulation_steps} is not supported \"\n f\"with fake teacher training. FT and Student share the same accelerator accumulation \"\n f\"counter, which causes gradient sync timing conflicts. Please set gradient_accumulation_steps=1.\"\n )\n\n # Determine model mode and initialize all components.\n # When FT is disabled (fake_teacher_config is None), default ft_use_lora to True\n # (value is irrelevant since FT branch will be skipped, but get_model_mode requires a bool).\n ft_use_lora = fake_teacher_config.use_lora if fake_teacher_config is not None else True\n \n model_mode = get_model_mode(fake_teacher_enabled, ft_use_lora, student_use_lora)\n logger.info(f\"Model mode: {model_mode.name}\", main_process_only=True)\n\n # Guard: LoRA mode is not supported with FSDP\n is_fsdp = accelerator.distributed_type == accelerator.distributed_type.FSDP\n if is_fsdp and (student_use_lora or ft_use_lora):\n raise ValueError(\n f\"LoRA mode is not supported with FSDP (student_use_lora={student_use_lora}, \"\n f\"ft_use_lora={ft_use_lora}). Please use DDP for LoRA training, or switch to \"\n f\"full fine-tuning mode for FSDP.\"\n )\n \n model_components = initialize_models(\n pipeline=pipeline, device=device, mode=model_mode,\n student_config=student_config, fake_teacher_config=fake_teacher_config,\n frozen_model_dtype=frozen_model_dtype,\n lora_target_modules=model_adapter.get_lora_target_modules(\n include_ff=student_config.lora_include_ff\n ),\n )\n \n transformer = model_components.transformer\n transformer_trainable_parameters = model_components.transformer_trainable_params\n fake_teacher = model_components.fake_teacher\n fake_teacher_trainable_parameters = model_components.fake_teacher_trainable_params\n \n # Log model setup info\n logger.info(f\"Fake Teacher initialized: mode={'LoRA' if is_ft_lora(model_mode) else 'Full'}\", main_process_only=True)\n \n # Gradient checkpointing for student\n if student_config.gradient_checkpointing:\n base_model = transformer.base_model.model if is_student_lora(model_mode) else transformer\n if hasattr(base_model, 'enable_gradient_checkpointing'):\n base_model.enable_gradient_checkpointing()\n logger.info(\"😊 [Memory优化] Student gradient checkpointing enabled\", main_process_only=True)\n \n # Gradient checkpointing for fake teacher (only when FT is enabled)\n if fake_teacher_enabled and fake_teacher_config.gradient_checkpointing \\\n and fake_teacher is not None and fake_teacher is not transformer:\n ft_base_model = fake_teacher.base_model.model if hasattr(fake_teacher, 'base_model') else fake_teacher\n if hasattr(ft_base_model, 'enable_gradient_checkpointing'):\n ft_base_model.enable_gradient_checkpointing()\n logger.info(\"😊 [Memory优化] Fake Teacher gradient checkpointing enabled\", main_process_only=True)\n\n # ==================== Training Statistics ====================\n train_batch_size = train_config.batch_size\n num_batches_per_epoch = config.sample.num_batches_per_epoch\n grad_accum = train_config.gradient_accumulation_steps\n \n # Student training steps\n student_update_ratio = student_config.update_ratio\n student_batches_per_epoch = num_batches_per_epoch\n student_steps_per_epoch = student_batches_per_epoch // grad_accum\n # Calculate total steps across all epochs, then divide by update_ratio\n # This correctly handles cases where steps_per_epoch < update_ratio\n student_total_raw_steps = student_steps_per_epoch * config.num_epochs\n total_steps = math.ceil(student_total_raw_steps / student_update_ratio) if student_update_ratio > 0 else 0\n student_effective_steps_per_epoch = total_steps / config.num_epochs if config.num_epochs > 0 else 0 # For logging only\n \n # Fake Teacher training steps (Fake Teacher trains every batch, no update_ratio)\n ft_batches_per_epoch = num_batches_per_epoch\n ft_steps_per_epoch = ft_batches_per_epoch // grad_accum\n ft_total_steps = ft_steps_per_epoch * config.num_epochs\n\n # ==================== Student Optimizer & LR Scheduler ====================\n adam_base_kwargs = dict(betas=(train_config.adam_beta1, train_config.adam_beta2),\n weight_decay=train_config.adam_weight_decay, eps=train_config.adam_epsilon)\n \n student_lr = student_config.learning_rate if is_student_lora(model_mode) else student_config.full_learning_rate\n student_adam_kwargs = dict(lr=student_lr, **adam_base_kwargs)\n if student_config.use_8bit_adam:\n try:\n import bitsandbytes as bnb\n optimizer = bnb.optim.AdamW8bit(transformer_trainable_parameters, **student_adam_kwargs)\n logger.info(\"😊 [Memory优化] Student using 8-bit AdamW\", main_process_only=True)\n except ImportError:\n optimizer = torch.optim.AdamW(transformer_trainable_parameters, **student_adam_kwargs)\n else:\n optimizer = torch.optim.AdamW(transformer_trainable_parameters, **student_adam_kwargs)\n\n student_warmup_steps = student_config.warmup_steps\n if student_warmup_steps == 0 and student_config.warmup_ratio > 0:\n student_warmup_steps = int(total_steps * student_config.warmup_ratio)\n student_lr_scheduler_type = student_config.lr_scheduler_type\n lr_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, create_lr_lambda(student_warmup_steps, total_steps, student_lr_scheduler_type))\n logger.info(f\"Student: lr={student_lr}, scheduler={student_lr_scheduler_type}, warmup={student_warmup_steps}, total_steps={total_steps}\", main_process_only=True)\n\n # ==================== Fake Teacher Optimizer & LR Scheduler ====================\n fake_teacher_optimizer, fake_teacher_lr_scheduler = None, None\n ft_step = 0\n ft_lr = None\n ft_lr_scheduler_type = None\n ft_warmup_steps = 0\n if fake_teacher_enabled:\n ft_lr = fake_teacher_config.learning_rate if is_ft_lora(model_mode) else fake_teacher_config.full_learning_rate\n ft_adam_kwargs = dict(lr=ft_lr, **adam_base_kwargs)\n\n if fake_teacher_config.use_8bit_adam:\n try:\n import bitsandbytes as bnb\n fake_teacher_optimizer = bnb.optim.AdamW8bit(fake_teacher_trainable_parameters, **ft_adam_kwargs)\n logger.info(\"😊 [Memory优化] Fake Teacher using 8-bit AdamW\", main_process_only=True)\n except ImportError:\n fake_teacher_optimizer = torch.optim.AdamW(fake_teacher_trainable_parameters, **ft_adam_kwargs)\n else:\n fake_teacher_optimizer = torch.optim.AdamW(fake_teacher_trainable_parameters, **ft_adam_kwargs)\n\n ft_warmup_steps = fake_teacher_config.warmup_steps\n if ft_warmup_steps == 0 and fake_teacher_config.warmup_ratio > 0:\n ft_warmup_steps = int(ft_total_steps * fake_teacher_config.warmup_ratio)\n ft_lr_scheduler_type = fake_teacher_config.lr_scheduler_type\n fake_teacher_lr_scheduler = torch.optim.lr_scheduler.LambdaLR(fake_teacher_optimizer, create_lr_lambda(ft_warmup_steps, ft_total_steps, ft_lr_scheduler_type))\n logger.info(f\"Fake Teacher: lr={ft_lr}, mode={'LoRA' if is_ft_lora(model_mode) else 'full'}, scheduler={ft_lr_scheduler_type}, warmup={ft_warmup_steps}, total_steps={ft_total_steps}\", main_process_only=True)\n\n # ==================== Datasets and Dataloaders ====================\n max_sample_samples = config.max_sample_samples\n max_eval_samples = config.eval.max_eval_samples\n sample_dataset_path = config.sample_dataset\n eval_dataset_path = config.eval_dataset\n \n sample_prompt_fn = config.prompt_fn\n eval_prompt_fn = config.eval_prompt_fn\n \n SampleDatasetClass = DATASET_REGISTRY.get(sample_prompt_fn)\n EvalDatasetClass = DATASET_REGISTRY.get(eval_prompt_fn)\n if not SampleDatasetClass:\n raise NotImplementedError(f\"Unknown prompt_fn '{sample_prompt_fn}'. Available: {list(DATASET_REGISTRY.keys())}\")\n if not EvalDatasetClass:\n raise NotImplementedError(f\"Unknown eval_prompt_fn '{eval_prompt_fn}'. Available: {list(DATASET_REGISTRY.keys())}\")\n\n # Build extra kwargs for datasets that need image-related parameters.\n # image_base_path may be empty string, in which case it falls back to None.\n image_size = (config.height, config.width)\n image_base_path = config.image_base_path or None\n image_kwargs = dict(image_size=image_size, image_base_path=image_base_path)\n image_dataset_types = {\"text_image_pair\", \"image_edit\"}\n sample_dataset_kwargs = image_kwargs if sample_prompt_fn in image_dataset_types else {}\n eval_dataset_kwargs = image_kwargs if eval_prompt_fn in image_dataset_types else {}\n\n sample_dataset = SampleDatasetClass(sample_dataset_path, \"train\", max_samples=max_sample_samples, **sample_dataset_kwargs)\n eval_dataset = EvalDatasetClass(eval_dataset_path, \"test\", max_samples=max_eval_samples, **eval_dataset_kwargs)\n logger.info(f\"Dataset sizes - sample: {len(sample_dataset)}, eval: {len(eval_dataset)}\", main_process_only=True)\n\n sample_sampler = DistributedKRepeatSampler(\n dataset=sample_dataset, batch_size=train_batch_size,\n k=config.sample.num_image_per_prompt, num_replicas=world_size, rank=rank, seed=config.seed,\n )\n sample_dataloader = DataLoader(sample_dataset, batch_sampler=sample_sampler, num_workers=0, collate_fn=sample_dataset.collate_fn, pin_memory=True)\n eval_sampler = DistributedSampler(eval_dataset, num_replicas=world_size, rank=rank, shuffle=False) if world_size > 1 else None\n eval_dataloader = DataLoader(eval_dataset, batch_size=config.eval.eval_batch_size, sampler=eval_sampler, collate_fn=eval_dataset.collate_fn, num_workers=0, pin_memory=True)\n\n executor = futures.ThreadPoolExecutor(max_workers=8)\n\n # ==================== Prepare with Accelerator ====================\n transformer, optimizer, lr_scheduler = accelerator.prepare(transformer, optimizer, lr_scheduler)\n patch_model_for_pipeline(transformer, accelerator.unwrap_model(transformer))\n pipeline.transformer = transformer\n transformer_trainable_parameters = [p for p in transformer.parameters() if p.requires_grad]\n\n # Prepare fake teacher with Accelerator.\n # For FT_LORA_STUDENT_LORA mode, FT shares the same transformer (handled by adapter switching).\n if fake_teacher_optimizer is not None and not model_mode == ModelMode.FT_LORA_STUDENT_LORA:\n fake_teacher, fake_teacher_optimizer, fake_teacher_lr_scheduler = accelerator.prepare(\n fake_teacher, fake_teacher_optimizer, fake_teacher_lr_scheduler\n )\n patch_model_for_pipeline(fake_teacher, accelerator.unwrap_model(fake_teacher))\n pipeline.fake_teacher = fake_teacher\n fake_teacher_trainable_parameters = [p for p in fake_teacher.parameters() if p.requires_grad]\n logger.info(\"😊 Fake Teacher prepared with Accelerator\", main_process_only=True)\n\n # Prepare real teacher with Accelerator (FSDP compatibility).\n real_teacher = pipeline.real_teacher # May be None for LoRA mode\n\n # ==================== LongCat ids cast proxy ====================\n if config.base_model == \"longcat\":\n pipeline.transformer = wrap_longcat_transformer_for_fsdp(pipeline.transformer)\n if hasattr(pipeline, \"fake_teacher\") and pipeline.fake_teacher is not None:\n pipeline.fake_teacher = wrap_longcat_transformer_for_fsdp(pipeline.fake_teacher)\n if real_teacher is not None:\n real_teacher = wrap_longcat_transformer_for_fsdp(real_teacher)\n pipeline.real_teacher = real_teacher\n logger.info(\"🐈 LongCat ids cast proxy applied to transformer/fake_teacher/real_teacher\",\n main_process_only=True)\n\n # ==================== Prompt Embeddings ====================\n # Pre-compute negative embeddings AFTER accelerator.prepare() so that\n neg_prompt_embed, neg_auxiliary_text_embed = model_adapter.compute_text_embeddings([\"\"], text_encoders, tokenizers, max_sequence_length=model_adapter.get_max_sequence_length(), device=device)\n neg_prompt_embeds = neg_prompt_embed.repeat(train_batch_size, 1, 1)\n neg_auxiliary_text_embeds = _repeat_embeds(neg_auxiliary_text_embed, train_batch_size)\n # Pre-compute negative model kwargs for CFG forward calls\n neg_model_kwargs = model_adapter.prepare_forward_kwargs(neg_prompt_embeds, neg_auxiliary_text_embeds)\n\n # ==================== Initialize Model Manager ====================\n # Update model_components with prepared (accelerator-wrapped) models.\n model_components = ModelComponents(\n transformer=transformer,\n transformer_trainable_params=transformer_trainable_parameters,\n fake_teacher=fake_teacher,\n fake_teacher_trainable_params=fake_teacher_trainable_parameters,\n real_teacher=real_teacher,\n mode=model_mode\n )\n \n model_manager = ModelManager(\n pipeline=pipeline,\n mode=model_mode,\n accelerator=accelerator,\n device=device,\n model_components=model_components,\n model_adapter=model_adapter,\n )\n \n logger.info(\"😊 ModelManager initialized\", main_process_only=True)\n\n logger.info(f\"***** Running training *****\", main_process_only=True)\n \n logger.info(f\"***** Training Config *****\", main_process_only=True)\n logger.info(f\" General: world_size={world_size}, epochs={config.num_epochs}, grad_accum={grad_accum}\", main_process_only=True)\n logger.info(f\" Train: per_gpu_bs={train_batch_size}, effective_bs={train_batch_size * world_size * grad_accum}, batches/epoch={num_batches_per_epoch}\", main_process_only=True)\n logger.info(f\" Eval: per_gpu_bs={config.eval.eval_batch_size}, effective_bs={config.eval.eval_batch_size * world_size}, freq={config.eval_freq} epochs\", main_process_only=True)\n logger.info(f\" Student: mode={'LoRA' if is_student_lora(model_mode) else 'Full'}, steps/epoch={student_steps_per_epoch}, update_ratio={student_update_ratio}, effective_steps/epoch={student_effective_steps_per_epoch:.2f}, total_steps={total_steps}, lr={student_lr}\", main_process_only=True)\n if fake_teacher_enabled:\n logger.info(f\" FakeTeacher: mode={'LoRA' if is_ft_lora(model_mode) else 'Full'}, steps/epoch={ft_steps_per_epoch}, total_steps={ft_total_steps}\", main_process_only=True)\n logger.info(f\" FakeTeacher: lr={ft_lr}, lr_scheduler={ft_lr_scheduler_type}, warmup={ft_warmup_steps}\", main_process_only=True)\n else:\n logger.info(\" FakeTeacher: disabled\", main_process_only=True)\n logger.info(f\" Output: {log_dir}\", main_process_only=True)\n\n eval_reward_config = config.eval.reward_fn or config.train.reward_fn\n eval_reward_fn = getattr(cdm.rewards.rewards, \"multi_score\")(device, eval_reward_config)\n\n # Offload reward models to CPU during training to save GPU memory.\n # They will be loaded back to GPU on-demand during eval phases.\n from cdm.rewards.rewards import get_scorer_models\n reward_scorer_models = get_scorer_models(eval_reward_fn)\n for score_name, scorer in reward_scorer_models:\n scorer.cpu()\n if hasattr(scorer, 'device'):\n scorer.device = torch.device('cpu')\n logger.info(f\"Offloaded reward model '{score_name}' to CPU to save GPU memory\", main_process_only=True)\n if reward_scorer_models:\n torch.cuda.empty_cache()\n\n # Resume from checkpoint\n first_epoch, student_step = 0, 0\n if config.resume_from:\n accelerator.load_state(config.resume_from)\n try:\n student_step = int(os.path.basename(config.resume_from).split(\"-\")[-1])\n except ValueError:\n student_step = 0\n logger.info(f\"Resumed from {config.resume_from}, student_step={student_step}\", main_process_only=True)\n\n ema_decay = student_config.ema_decay\n ema_device = device # Keep EMA on GPU for faster updates (no cross-device copy)\n if config.train.student.ema:\n # FSDP2: EMA operates on local shards via _get_data() in EMAModuleWrapper\n # DDP/single-GPU: operates on full parameters directly\n ema = EMAModuleWrapper(transformer_trainable_parameters, decay=ema_decay, update_step_interval=1, device=ema_device)\n logger.info(f\"😊 EMA parameters stored on GPU (uses ~{sum(p.numel() * p.element_size() for p in transformer_trainable_parameters) / 1024**3:.1f} GB GPU memory)\", main_process_only=True)\n else:\n ema = None\n\n # Load EMA state from checkpoint if available\n if config.resume_from and ema is not None:\n ema_path = os.path.join(config.resume_from, \"ema_state.pt\")\n if os.path.exists(ema_path):\n ema.load_state_dict(torch.load(ema_path, map_location=device))\n logger.info(\"Loaded EMA state from checkpoint\", main_process_only=True)\n else:\n logger.warning(\"No EMA state found in checkpoint, EMA initialized from current model\", main_process_only=True)\n\n sample_iter = iter(sample_dataloader)\n optimizer.zero_grad()\n\n # Note: old model EMA update is now performed after each student step in the training loop\n\n train_start_time = time.time()\n\n # ==================== Config node aliases (for concise access in the loop) ====================\n sample_config = config.sample\n loss_config = config.train.loss\n cfg_loss = loss_config.cfg\n ddm_loss = loss_config.ddm\n cdm_loss = loss_config.cdm\n # FT alias is only valid when fake_teacher is enabled.\n ft = fake_teacher_config if fake_teacher_enabled else None\n\n # Global sigma truncation bounds\n sigma_min = train_config.sigma_min\n sigma_max = train_config.sigma_max\n\n if fake_teacher_enabled and ft.ema_from_student:\n logger.info(f\"FT←Student EMA merge: enabled, decay={ft.ema_from_student_decay}, interval={ft.ema_from_student_interval}\", main_process_only=True)\n \n# Custom sigmas for denoising schedule (used by both training sampling and eval).\n # Set config.custom_sigmas to None to use the default schedule from num_inference_steps.\n custom_sigmas = list(config.custom_sigmas) if config.custom_sigmas is not None else None\n if custom_sigmas is not None:\n logger.info(f\"Custom sigmas: {custom_sigmas}\", main_process_only=True)\n \n # Random sigma schedule: randomize sigmas each epoch for ODE discretization generalization\n if sample_config.random_sigma_schedule:\n logger.info(f\"Random sigma schedule: enabled, steps=[{sample_config.random_sigma_num_steps_min}, {sample_config.random_sigma_num_steps_max}], \"\n f\"sigma_range=[{sample_config.random_sigma_min}, {sample_config.random_sigma_max}]\", main_process_only=True)\n \n if cdm_loss.teacher_x0_base == \"student_pred\":\n assert cdm_loss.teacher_xt_mode != \"share_student\", (\n \"cdm.teacher_x0_base='student_pred' is incompatible with cdm.teacher_xt_mode='share_student' \"\n \"(share_student forces teacher to reuse student's xt, leaving no room for a re-noised x0_pred). \"\n \"Use 'fresh_shared_t' or 'fresh_independent' instead.\"\n )\n if cdm_loss.teacher_xt_mode == \"share_student\":\n assert cdm_loss.teacher_x0_base == cdm_loss.student_x0_base, (\n f\"cdm.teacher_xt_mode='share_student' requires teacher_x0_base ({cdm_loss.teacher_x0_base!r}) \"\n f\"to equal student_x0_base ({cdm_loss.student_x0_base!r}).\"\n )\n logger.info(f\"CDM xt mode: student={cdm_loss.student_xt_mode}, teacher={cdm_loss.teacher_xt_mode}, \"\n f\"student_x0_base={cdm_loss.student_x0_base}, teacher_x0_base={cdm_loss.teacher_x0_base}\",\n main_process_only=True)\n\n logger.info(f\"Loss: cfg={cfg_loss.enabled}(w={cfg_loss.weight}), ddm={ddm_loss.enabled}(w={ddm_loss.weight}), cdm={cdm_loss.enabled}(w={cdm_loss.weight})\", main_process_only=True)\n\n\n epoch_pbar = tqdm(range(first_epoch, config.num_epochs), desc=\"Training Progress\", \n position=0, disable=not accelerator.is_main_process)\n \n for epoch in epoch_pbar:\n epoch_pbar.set_description(f\"Epoch {epoch}/{config.num_epochs - 1}\")\n \n # --- Merged Training Loop Body ---\n torch.cuda.reset_peak_memory_stats(device)\n \n unwrapped_transformer = accelerator.unwrap_model(transformer)\n model_dtype = next(unwrapped_transformer.parameters()).dtype\n \n # ========== Evaluation Phase ==========\n should_eval = epoch % config.eval_freq == 0\n if config.eval.skip_first_eval and epoch == 0:\n should_eval = False\n \n if should_eval:\n if config.eval.distillation_enabled:\n eval_distillation_fn(pipeline, eval_dataloader, text_encoders, tokenizers, config,\n device, rank, world_size, student_step, eval_reward_fn, executor,\n mixed_precision_dtype, ema, transformer_trainable_parameters, exp_logger, transformer_ddp=transformer, model_adapter=model_adapter, accelerator=accelerator)\n # ========== Fake Teacher Visualization ==========\n ft_viz_config = config.eval.fake_teacher_viz\n ft_viz_enabled = ft_viz_config is not None and ft_viz_config.enabled\n should_viz_fake_teacher = ft_viz_enabled\n \n if should_viz_fake_teacher:\n eval_fake_teacher_viz_fn(\n pipeline=pipeline, test_dataloader=eval_dataloader, text_encoders=text_encoders, tokenizers=tokenizers,\n config=config, device=device, rank=rank, world_size=world_size, global_step=student_step, epoch=epoch,\n mixed_precision_dtype=mixed_precision_dtype, exp_logger=exp_logger, accelerator=accelerator,\n transformer_ddp=transformer, model_adapter=model_adapter,\n )\n\n # ========== Checkpoint Saving ==========\n if epoch % config.save_freq == 0 and epoch > 0:\n checkpoint_save_path = os.path.join(checkpoint_dir, f\"checkpoint-{student_step}\")\n if accelerator.is_main_process:\n os.makedirs(checkpoint_save_path, exist_ok=True)\n accelerator.wait_for_everyone()\n accelerator.save_state(checkpoint_save_path)\n if ema is not None and accelerator.is_main_process:\n torch.save(ema.state_dict(), os.path.join(checkpoint_save_path, \"ema_state.pt\"))\n # FSDP: save full (unsharded) state dicts for prepare_student_pipeline compatibility\n save_fsdp_full_checkpoint(accelerator, transformer, ema, checkpoint_save_path, logger)\n if accelerator.is_main_process:\n logger.info(f\"Saved checkpoint to {checkpoint_save_path}\")\n\n if accelerator.is_main_process and exp_logger is not None:\n exp_logger.add_scalar(\"global/epoch\", epoch, student_step)\n\n # ========== Unified Sample-Train Loop ==========\n # Each batch: sample → FT train → Student train (no collect-then-split)\n ft_model_for_training = transformer if model_mode == ModelMode.FT_LORA_STUDENT_LORA else fake_teacher\n gradient_update_times = 0\n info_accumulated = defaultdict(list)\n\n batch_pbar = tqdm(range(num_batches_per_epoch), desc=f\"Epoch {epoch}\",\n position=0, disable=not accelerator.is_main_process)\n\n for batch_idx in batch_pbar:\n # ---- 1. Sampling ----\n unwrapped_transformer.eval()\n if hasattr(sample_sampler, \"set_epoch\") and isinstance(sample_sampler, DistributedKRepeatSampler):\n sample_sampler.set_epoch(epoch * num_batches_per_epoch + batch_idx)\n\n prompts, prompt_metadata = next(sample_iter)\n\n num_steps = config.sample.num_steps\n model_manager.prepare_for_sampling(sample_config.sampling_adapter)\n current_batch_size = len(prompts)\n\n with create_autocast_context(accelerator), torch.no_grad():\n pipeline_kwargs = dict(\n prompt=prompts,\n num_inference_steps=num_steps, guidance_scale=config.sample.guidance_scale,\n height=config.height, width=config.width,\n deterministic=config.sample.deterministic,\n )\n\n # Random sigma schedule: generate new random sigmas each epoch.\n if sample_config.random_sigma_schedule:\n import random as _random\n _len_rng = _random.Random((epoch, batch_idx))\n rand_num_steps = _len_rng.randint(sample_config.random_sigma_num_steps_min, sample_config.random_sigma_num_steps_max)\n inner_sigmas = sorted(\n [_random.uniform(sample_config.random_sigma_min, sample_config.random_sigma_max) for _ in range(rand_num_steps - 1)],\n reverse=True,\n )\n epoch_sigmas = [1.0] + inner_sigmas\n pipeline_kwargs[\"sigmas\"] = epoch_sigmas\n pipeline_kwargs[\"num_inference_steps\"] = rand_num_steps\n num_steps = rand_num_steps\n # Custom sigmas override the default schedule from num_inference_steps\n elif custom_sigmas is not None:\n pipeline_kwargs[\"sigmas\"] = custom_sigmas\n\n # LongCat: pass enable_prompt_rewrite for training sampling\n if config.base_model == \"longcat\":\n pipeline_kwargs[\"enable_prompt_rewrite\"] = config.eval.enable_prompt_rewrite\n\n # backward simulation\n (_, all_xt, all_x0, all_v, prompt_embeds, auxiliary_text_embeds,\n latent_image_ids) = pipeline_infer(pipeline, return_images=False, **pipeline_kwargs)\n\n # Pre-generate shared trajectory indices for this batch when enabled.\n if train_config.share_latent_indices:\n shared_latent_indices_cpu = torch.randint(\n 0, num_steps, (current_batch_size,), device=device).cpu().tolist()\n else:\n shared_latent_indices_cpu = None\n\n # ---- 2. Fake Teacher Training ----\n model_manager.prepare_for_ft_training()\n\n # Extract FT target x0 from trajectory based on ft.x0_source config\n # FT samples its own trajectory indices independently (only when needed)\n if ft.x0_source == \"x0_final\":\n ft_batch_x0 = all_xt[-1].to(model_dtype)\n else:\n _, ft_indices_cpu = make_indices(\n shared_latent_indices_cpu, current_batch_size, num_steps, device)\n if ft.x0_source == \"x0_next\":\n ft_batch_x0 = extract_traj(all_x0, [idx + 1 for idx in ft_indices_cpu]).to(model_dtype)\n else: # x0_current\n ft_batch_x0 = extract_traj(all_x0, ft_indices_cpu).to(model_dtype)\n ft_batch_prompt_embeds = prompt_embeds.to(model_dtype)\n ft_batch_auxiliary_embeds = auxiliary_text_embeds if auxiliary_text_embeds is not None else None\n ft_current_batch_size = ft_batch_x0.shape[0]\n\n # Note: latent_image_ids is global (per-pipeline-call), independent of trajectory indices.\n ft_model_kwargs = model_adapter.prepare_forward_kwargs(\n ft_batch_prompt_embeds, ft_batch_auxiliary_embeds, img_ids=latent_image_ids)\n\n ft_sigma = sample_sigmas(\n ft_current_batch_size, ft.sigma_sampling_method, device,\n sigma_min=sigma_min, sigma_max=sigma_max,\n logit_mean=ft.logit_mean, logit_std=ft.logit_std)\n ft_sigma_expanded = ft_sigma.view(-1, *([1] * (len(ft_batch_x0.shape) - 1)))\n ft_timestep = (ft_sigma * 1000.0).float()\n\n ft_noise = torch.randn_like(ft_batch_x0)\n ft_xt = (1 - ft_sigma_expanded) * ft_batch_x0.detach() + ft_sigma_expanded * ft_noise\n\n ft_xt = ft_xt.to(model_dtype)\n with accelerator.accumulate(ft_model_for_training):\n with create_autocast_context(accelerator):\n ft_prediction = model_manager.fake_teacher_forward(\n ft_xt, ft_timestep, requires_grad=True, **ft_model_kwargs)\n ft_diffusion_loss, ft_diffusion_info = compute_fake_teacher_diffusion_loss(\n fake_teacher_prediction=ft_prediction, x0=ft_batch_x0.detach(), noise=ft_noise, t_expanded=ft_sigma_expanded)\n scaled_loss = loss_config.diffusion.weight * ft_diffusion_loss\n\n accelerator.backward(scaled_loss)\n\n if accelerator.sync_gradients:\n if fake_teacher_optimizer:\n ft_grad_stats = compute_grad_stats(parameters=fake_teacher_trainable_parameters)\n if accelerator.distributed_type == DistributedType.FSDP:\n torch.nn.utils.clip_grad_norm_(fake_teacher_trainable_parameters, config.train.max_grad_norm)\n else:\n accelerator.clip_grad_norm_(fake_teacher_trainable_parameters, config.train.max_grad_norm)\n fake_teacher_optimizer.step()\n if fake_teacher_lr_scheduler is not None:\n fake_teacher_lr_scheduler.step()\n fake_teacher_optimizer.zero_grad()\n\n # FT ← Student EMA merge: blend student params into FT after gradient update\n if ft.ema_from_student and ft_step % ft.ema_from_student_interval == 0:\n ema_merge_ft_with_student(\n ft_params=fake_teacher_trainable_parameters,\n student_params=transformer_trainable_parameters,\n decay=ft.ema_from_student_decay,\n )\n\n ft_step += 1\n\n if accelerator.is_main_process and exp_logger is not None:\n current_ft_lr = fake_teacher_lr_scheduler.get_last_lr()[0] if fake_teacher_lr_scheduler is not None else ft_lr\n log_fake_teacher_metrics(exp_logger, ft_step, ft_diffusion_loss, ft_sigma,\n ft_diffusion_info, ft_grad_stats, current_ft_lr, ft_current_batch_size)\n\n del ft_prediction, ft_xt, ft_noise\n\n model_manager.finish_ft_training()\n batch_pbar.set_postfix({\"ft_loss\": f\"{ft_diffusion_loss.item():.4f}\", \"σ\": f\"{ft_sigma.mean().item():.3f}\", \"ft_step\": ft_step})\n\n # ---- 3. Student Training ----\n should_train_student = (student_update_ratio == 1) or (ft_step > 0 and ft_step % student_update_ratio == 0)\n if not should_train_student:\n continue\n\n model_manager.prepare_for_student_training()\n\n batch_prompt_embeds = prompt_embeds.to(model_dtype)\n batch_auxiliary_embeds = auxiliary_text_embeds if auxiliary_text_embeds is not None else None\n batch_model_kwargs = model_adapter.prepare_forward_kwargs(\n batch_prompt_embeds, batch_auxiliary_embeds, img_ids=latent_image_ids)\n\n # Shared with CDM (used as x0_final base when cdm_*_x0_base = \"x0_final\"). Cheap, no indexing.\n batch_x0_final = all_xt[-1].to(model_dtype)\n\n current_batch_size = len(batch_prompt_embeds)\n # Build negative model kwargs for CFG\n batch_neg_model_kwargs = dict(neg_model_kwargs)\n # Inject latent_image_ids into neg kwargs (img_ids is the same for cond/uncond)\n if latent_image_ids is not None:\n batch_neg_model_kwargs[\"img_ids\"] = latent_image_ids\n\n with accelerator.accumulate(transformer):\n loss_terms = {}\n zero_tensor = torch.tensor(0.0, device=device)\n total_loss = torch.tensor(0.0, device=device)\n\n # ========== Loss 1: CFG on_trajectory ==========\n if cfg_loss.enabled and cfg_loss.weight != 0:\n with create_autocast_context(accelerator):\n # Sample trajectory indices for CFG (shared across branches when enabled)\n cfg_latent_indices, cfg_latent_indices_cpu = make_indices(\n shared_latent_indices_cpu, current_batch_size, num_steps, device)\n cfg_batch_xt = extract_traj(all_xt, cfg_latent_indices_cpu).float()\n cfg_batch_stop_timesteps = pipeline.scheduler.timesteps[cfg_latent_indices.cpu()].to(device).float()\n cfg_stop_sigma_base = (cfg_batch_stop_timesteps / 1000.0).float()\n cfg_stop_sigma_base_expanded = cfg_stop_sigma_base.view(-1, *([1] * (len(cfg_batch_xt.shape) - 1)))\n\n cfg_student_sigma = cfg_stop_sigma_base\n cfg_student_xt = cfg_batch_xt\n cfg_student_sigma_expanded = cfg_stop_sigma_base_expanded\n cfg_student_timesteps = cfg_batch_stop_timesteps\n \n cfg_student_prediction = model_manager.student_forward(\n cfg_student_xt.to(model_dtype), cfg_student_timesteps, \n **batch_model_kwargs)\n \n cfg_x0_pred = cfg_student_xt - cfg_student_sigma_expanded * cfg_student_prediction\n \n # Teacher re-noise base: always use student's online x0 prediction.\n cfg_x0_for_teacher = cfg_x0_pred\n # Teacher sigma upper bound: always the student's stop sigma.\n cfg_teacher_upper_bound = cfg_student_sigma\n \n cfg_teacher_sigma = sample_sigmas(\n current_batch_size, cfg_loss.teacher_sigma_method, device,\n sigma_min=sigma_min, sigma_max=sigma_max,\n logit_mean=cfg_loss.teacher_sigma_logit_mean, logit_std=cfg_loss.teacher_sigma_logit_std,\n upper_bound=cfg_teacher_upper_bound)\n cfg_teacher_sigma_expanded = cfg_teacher_sigma.view(-1, *([1] * (len(cfg_x0_for_teacher.shape) - 1)))\n cfg_teacher_timesteps = (cfg_teacher_sigma * 1000.0).float()\n cfg_teacher_noise = torch.randn_like(cfg_x0_for_teacher)\n cfg_teacher_xt = (1 - cfg_teacher_sigma_expanded) * cfg_x0_for_teacher.detach() + cfg_teacher_sigma_expanded * cfg_teacher_noise\n \n cfg_uncond_pred, cfg_cond_pred = model_manager.teacher_forward_cfg_batched(\n cfg_teacher_xt.to(model_dtype), cfg_teacher_timesteps,\n batch_model_kwargs, batch_neg_model_kwargs)\n \n cfg_loss_val, cfg_loss_info = compute_cfg_loss(\n student_prediction=cfg_student_prediction,\n real_teacher_cond_prediction=cfg_cond_pred,\n real_teacher_uncond_prediction=cfg_uncond_pred,\n xt=cfg_student_xt, xt_teacher=cfg_teacher_xt,\n student_sigma_expanded=cfg_student_sigma_expanded,\n teacher_sigma_expanded=cfg_teacher_sigma_expanded,\n guidance_scale=cfg_loss.guidance_scale,\n return_per_sample=False,\n norm_clip_min=cfg_loss.norm_clip_min)\n \n cfg_loss_weighted = cfg_loss.weight * cfg_loss_val\n total_loss = total_loss + cfg_loss_weighted\n loss_terms.update({f\"cfg/{k}\": v for k, v in cfg_loss_info.items()})\n loss_terms.update({\"loss/cfg_raw\": cfg_loss_val.detach(), \"loss/cfg\": cfg_loss_weighted.detach()})\n\n else:\n loss_terms[\"loss/cfg\"] = zero_tensor\n\n # ========== Loss 2: DM on_trajectory ==========\n if ddm_loss.enabled and ddm_loss.weight != 0:\n with create_autocast_context(accelerator):\n # Sample trajectory indices for DDM (shared across branches when enabled)\n ddm_latent_indices, ddm_latent_indices_cpu = make_indices(\n shared_latent_indices_cpu, current_batch_size, num_steps, device)\n # Extract DDM trajectory data on-demand\n ddm_batch_xt = extract_traj(all_xt, ddm_latent_indices_cpu)\n ddm_batch_stop_timesteps = pipeline.scheduler.timesteps[ddm_latent_indices.cpu()].to(device)\n ddm_student_timesteps = ddm_batch_stop_timesteps.float()\n ddm_stop_sigma = (ddm_student_timesteps / 1000.0).float()\n \n ddm_student_sigma = ddm_stop_sigma\n ddm_student_xt = ddm_batch_xt.float().to(model_dtype)\n ddm_student_sigma_expanded = ddm_stop_sigma.view(-1, *([1] * (len(ddm_student_xt.shape) - 1)))\n \n ddm_student_prediction = model_manager.student_forward(\n ddm_student_xt, ddm_student_timesteps, **batch_model_kwargs)\n \n ddm_x0_pred = ddm_student_xt - ddm_student_sigma_expanded * ddm_student_prediction\n \n # Teacher re-noise base: always use student's online x0 prediction.\n ddm_x0_for_teacher = ddm_x0_pred\n # Teacher sigma upper bound: always the student's stop sigma.\n ddm_teacher_upper_bound = ddm_student_sigma\n \n ddm_teacher_sigma = sample_sigmas(\n current_batch_size, ddm_loss.teacher_sigma_method, device,\n sigma_min=sigma_min, sigma_max=sigma_max,\n logit_mean=ddm_loss.teacher_sigma_logit_mean, logit_std=ddm_loss.teacher_sigma_logit_std,\n upper_bound=ddm_teacher_upper_bound)\n ddm_teacher_sigma_expanded = ddm_teacher_sigma.view(-1, *([1] * (len(ddm_x0_for_teacher.shape) - 1)))\n ddm_teacher_timesteps = (ddm_teacher_sigma * 1000.0).float()\n ddm_teacher_noise = torch.randn_like(ddm_x0_for_teacher)\n ddm_teacher_xt = (1 - ddm_teacher_sigma_expanded) * ddm_x0_for_teacher.detach() + ddm_teacher_sigma_expanded * ddm_teacher_noise\n \n ddm_real_teacher_prediction = model_manager.teacher_forward(\n ddm_teacher_xt.to(model_dtype), ddm_teacher_timesteps,\n **batch_model_kwargs)\n ddm_fake_teacher_prediction = model_manager.fake_teacher_forward(\n ddm_teacher_xt.to(model_dtype), ddm_teacher_timesteps,\n requires_grad=False, **batch_model_kwargs)\n \n # Get CFG gradient direction for alignment analysis (if CFG loss was computed)\n ddm_cfg_grad_direction = cfg_loss_info.get(\"cfg_grad_direction\", None) if cfg_loss.enabled and cfg_loss.weight != 0 else None\n \n ddm_loss_val, ddm_loss_info = compute_dm_loss(\n student_prediction=ddm_student_prediction,\n real_teacher_prediction=ddm_real_teacher_prediction,\n fake_teacher_prediction=ddm_fake_teacher_prediction,\n xt=ddm_student_xt, xt_teacher=ddm_teacher_xt,\n student_sigma_expanded=ddm_student_sigma_expanded,\n teacher_sigma_expanded=ddm_teacher_sigma_expanded,\n return_per_sample=False,\n cfg_grad_direction=ddm_cfg_grad_direction,\n norm_clip_min=ddm_loss.norm_clip_min)\n \n ddm_loss_weighted = ddm_loss.weight * ddm_loss_val\n total_loss = total_loss + ddm_loss_weighted\n loss_terms.update({f\"ddm/{k}\": v for k, v in ddm_loss_info.items()})\n loss_terms.update({\"loss/ddm_raw\": ddm_loss_val.detach(), \"loss/ddm\": ddm_loss_weighted.detach()})\n\n else:\n loss_terms[\"loss/ddm\"] = zero_tensor\n\n # ========== Loss 3: DM fresh_noise ==========\n if cdm_loss.enabled and cdm_loss.weight != 0:\n with create_autocast_context(accelerator):\n # ---------- Step A: trajectory indices (if needed) and student x0_base tensor ----------\n if cdm_loss.student_xt_mode in (\"trajectory_eps\", \"trajectory_eps_last\", \"trajectory_xt\"):\n if cdm_loss.student_xt_mode == \"trajectory_eps_last\":\n # Always pick the last trajectory step (closest to x0) instead of random sampling.\n # This branch is intentionally NOT shared (it is deterministic by design).\n cdm_traj_indices_cpu = [num_steps - 1] * current_batch_size\n else:\n _, cdm_traj_indices_cpu = make_indices(\n shared_latent_indices_cpu, current_batch_size, num_steps, device)\n if cdm_loss.student_x0_base == \"x0_next\":\n cdm_student_x0_tensor = extract_traj(\n all_x0, [i + 1 for i in cdm_traj_indices_cpu]).to(model_dtype)\n else: # x0_final\n cdm_student_x0_tensor = all_xt[-1].to(model_dtype)\n else:\n # \"fresh\" mode: independently sample x0_next if needed, else use x0_final\n if cdm_loss.student_x0_base == \"x0_next\":\n _, cdm_x0_next_indices_cpu = make_indices(\n shared_latent_indices_cpu, current_batch_size, num_steps, device)\n cdm_student_x0_tensor = extract_traj(\n all_x0, [i + 1 for i in cdm_x0_next_indices_cpu]).to(model_dtype)\n else:\n cdm_student_x0_tensor = batch_x0_final\n\n # ---------- Step B: build student (sigma, xt) by mode ----------\n if cdm_loss.student_xt_mode == \"trajectory_xt\":\n # Directly reuse the on-trajectory point: (traj_xt, σ_traj)\n cdm_student_sigma = torch.tensor(\n [pipeline.scheduler.sigmas[i].item() for i in cdm_traj_indices_cpu],\n device=device, dtype=torch.float32)\n cdm_student_xt = extract_traj(all_xt, cdm_traj_indices_cpu)\n elif cdm_loss.student_xt_mode in (\"trajectory_eps\", \"trajectory_eps_last\"):\n # Reuse on-trajectory noise direction (v); independently sampled σ_s with v-extrapolation:\n # xt = traj_xt + (σ_s - σ_traj) * traj_v\n # For \"trajectory_eps_last\", traj index is fixed to the last step (closest to x0).\n cdm_traj_xt = extract_traj(all_xt, cdm_traj_indices_cpu)\n cdm_traj_v = extract_traj(all_v, cdm_traj_indices_cpu)\n cdm_traj_stop_sigma = torch.tensor(\n [pipeline.scheduler.sigmas[i].item() for i in cdm_traj_indices_cpu],\n device=device, dtype=torch.float32)\n cdm_student_sigma = sample_sigmas(\n current_batch_size, cdm_loss.student_sigma_method, device,\n sigma_min=sigma_min, sigma_max=sigma_max,\n logit_mean=cdm_loss.student_sigma_logit_mean,\n logit_std=cdm_loss.student_sigma_logit_std,\n upper_bound=cdm_traj_stop_sigma)\n sigma_diff = (cdm_student_sigma - cdm_traj_stop_sigma).view(\n -1, *([1] * (cdm_traj_xt.ndim - 1)))\n cdm_student_xt = cdm_traj_xt + sigma_diff * cdm_traj_v\n else: # \"fresh\"\n # Off-trajectory: xt = (1 - σ_s) * x0_base + σ_s * ε_fresh\n cdm_student_sigma = sample_sigmas(\n current_batch_size, cdm_loss.student_sigma_method, device,\n sigma_min=sigma_min, sigma_max=sigma_max,\n logit_mean=cdm_loss.student_sigma_logit_mean,\n logit_std=cdm_loss.student_sigma_logit_std)\n sigma_exp = cdm_student_sigma.view(-1, *([1] * (cdm_student_x0_tensor.ndim - 1)))\n cdm_noise = torch.randn_like(cdm_student_x0_tensor)\n cdm_student_xt = (1 - sigma_exp) * cdm_student_x0_tensor.detach() + sigma_exp * cdm_noise\n\n # ---------- Step C: derive expanded sigma & timesteps ----------\n cdm_student_sigma_expanded = cdm_student_sigma.view(\n -1, *([1] * (cdm_student_xt.ndim - 1)))\n cdm_student_timesteps = (cdm_student_sigma * 1000.0).float()\n\n # ---------- Step E: student forward (moved earlier so teacher can consume student's x0_pred) ----------\n cdm_student_prediction = model_manager.student_forward(\n cdm_student_xt.to(model_dtype), cdm_student_timesteps,\n **batch_model_kwargs)\n\n # ---------- Step A2: teacher x0_base tensor (deferred until student_forward) ----------\n # When teacher_xt_mode == \"share_student\", teacher must reuse student's x0 (no decoupling possible;\n # already validated above that teacher_x0_base == student_x0_base in this case).\n if cdm_loss.teacher_xt_mode == \"share_student\":\n cdm_teacher_x0_tensor = cdm_student_x0_tensor\n elif cdm_loss.teacher_x0_base == \"student_pred\":\n # Use student's online x0 prediction as the teacher's x0 base.\n # Detach to prevent gradients flowing back through the teacher's noised input.\n cdm_teacher_x0_tensor = (\n cdm_student_xt - cdm_student_sigma_expanded * cdm_student_prediction\n ).detach().to(model_dtype)\n elif cdm_loss.teacher_x0_base == \"x0_next\":\n # Sample trajectory indices for teacher's x0_next (shared across branches when enabled)\n _, teacher_traj_indices_cpu = make_indices(\n shared_latent_indices_cpu, current_batch_size, num_steps, device)\n cdm_teacher_x0_tensor = extract_traj(\n all_x0, [i + 1 for i in teacher_traj_indices_cpu]).to(model_dtype)\n else: # x0_final\n cdm_teacher_x0_tensor = all_xt[-1].to(model_dtype)\n\n # ---------- Step D: build teacher (sigma, xt) by mode ----------\n if cdm_loss.teacher_xt_mode == \"share_student\":\n # Teacher sees exactly the same (xt, t) as student\n cdm_teacher_sigma = cdm_student_sigma\n cdm_teacher_sigma_expanded = cdm_student_sigma_expanded\n cdm_teacher_timesteps = cdm_student_timesteps\n cdm_teacher_xt = cdm_student_xt\n elif cdm_loss.teacher_xt_mode == \"fresh_shared_t\":\n # Same t as student, but xt rebuilt from x0_base + fresh noise (decouples noise)\n cdm_teacher_sigma = cdm_student_sigma\n cdm_teacher_sigma_expanded = cdm_student_sigma_expanded\n cdm_teacher_timesteps = cdm_student_timesteps\n cdm_teacher_noise = torch.randn_like(cdm_teacher_x0_tensor)\n cdm_teacher_xt = ((1 - cdm_teacher_sigma_expanded) * cdm_teacher_x0_tensor.detach()\n + cdm_teacher_sigma_expanded * cdm_teacher_noise)\n else: # \"fresh_independent\"\n # Independently sampled σ_t and ε_fresh (strongest decoupling)\n cdm_teacher_sigma = sample_sigmas(\n current_batch_size, cdm_loss.teacher_sigma_method, device,\n sigma_min=sigma_min, sigma_max=sigma_max,\n logit_mean=cdm_loss.teacher_sigma_logit_mean,\n logit_std=cdm_loss.teacher_sigma_logit_std)\n cdm_teacher_sigma_expanded = cdm_teacher_sigma.view(\n -1, *([1] * (cdm_teacher_x0_tensor.ndim - 1)))\n cdm_teacher_timesteps = (cdm_teacher_sigma * 1000.0).float()\n cdm_teacher_noise = torch.randn_like(cdm_teacher_x0_tensor)\n cdm_teacher_xt = ((1 - cdm_teacher_sigma_expanded) * cdm_teacher_x0_tensor.detach()\n + cdm_teacher_sigma_expanded * cdm_teacher_noise)\n\n # DM gradient mode: use real - fake teacher gradient (DMD2 style)\n cdm_real_teacher_prediction = model_manager.teacher_forward(\n cdm_teacher_xt.to(model_dtype), cdm_teacher_timesteps,\n **batch_model_kwargs)\n cdm_fake_teacher_prediction = model_manager.fake_teacher_forward(\n cdm_teacher_xt.to(model_dtype), cdm_teacher_timesteps,\n requires_grad=False, **batch_model_kwargs)\n\n # Get CFG gradient direction for alignment analysis (if CFG loss was computed)\n cdm_cfg_grad_direction = cfg_loss_info.get(\"cfg_grad_direction\", None) if cfg_loss.enabled and cfg_loss.weight != 0 else None\n\n cdm_loss_val, cdm_loss_info = compute_dm_loss(\n student_prediction=cdm_student_prediction,\n real_teacher_prediction=cdm_real_teacher_prediction,\n fake_teacher_prediction=cdm_fake_teacher_prediction,\n xt=cdm_student_xt, xt_teacher=cdm_teacher_xt,\n student_sigma_expanded=cdm_student_sigma_expanded,\n teacher_sigma_expanded=cdm_teacher_sigma_expanded,\n return_per_sample=False,\n cfg_grad_direction=cdm_cfg_grad_direction,\n norm_clip_min=cdm_loss.norm_clip_min)\n\n cdm_loss_weighted = cdm_loss.weight * cdm_loss_val\n total_loss = total_loss + cdm_loss_weighted\n loss_terms.update({f\"cdm/{k}\": v for k, v in cdm_loss_info.items()})\n loss_terms.update({\"loss/cdm_raw\": cdm_loss_val.detach(), \"loss/cdm\": cdm_loss_weighted.detach()})\n loss_terms[\"cdm/student_sigma_mean\"] = cdm_student_sigma.mean().detach()\n else:\n loss_terms[\"loss/cdm\"] = zero_tensor\n\n # ========== Unified backward ==========\n accelerator.backward(total_loss)\n\n loss_terms.update({\n \"weights/cfg\": torch.tensor(cfg_loss.weight, device=device),\n \"weights/ddm\": torch.tensor(ddm_loss.weight, device=device),\n \"weights/cdm\": torch.tensor(cdm_loss.weight, device=device),\n \"loss/total\": total_loss.detach(),\n })\n\n for k, v in loss_terms.items():\n info_accumulated[k].append(v)\n\n pbar_info = {\"ft_step\": ft_step, \"stu_step\": student_step}\n for loss_key in [\"loss/total\", \"loss/cfg\", \"loss/ddm\", \"loss/cdm\"]:\n if loss_key in loss_terms:\n short_key = loss_key.split(\"/\")[-1]\n pbar_info[short_key] = f\"{loss_terms[loss_key].item():.4f}\"\n batch_pbar.set_postfix(pbar_info)\n\n del loss_terms\n if accelerator.sync_gradients:\n student_grad_stats = compute_grad_stats(parameters=transformer_trainable_parameters)\n if accelerator.distributed_type == DistributedType.FSDP:\n torch.nn.utils.clip_grad_norm_(transformer_trainable_parameters, config.train.max_grad_norm)\n else:\n accelerator.clip_grad_norm_(transformer.parameters(), config.train.max_grad_norm)\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n gradient_update_times += 1\n\n log_info = {k: torch.mean(torch.stack(v_list)).item() for k, v_list in info_accumulated.items()}\n info_tensor = torch.tensor([log_info[k] for k in sorted(log_info.keys())], device=device)\n info_tensor = accelerator.reduce(info_tensor, reduction=\"mean\")\n reduced_log_info = {k: info_tensor[i].item() for i, k in enumerate(sorted(log_info.keys()))}\n \n if accelerator.is_main_process and exp_logger is not None:\n exp_logger.add_scalar(\"student/step\", student_step, student_step)\n exp_logger.add_scalar(\"student/learning_rate\", lr_scheduler.get_last_lr()[0], student_step)\n exp_logger.add_scalar(\"student/batch_size\", train_batch_size, student_step)\n exp_logger.add_scalar(\"student/gradient_update_times\", gradient_update_times, student_step)\n exp_logger.add_scalar(\"student/update_ratio\", student_update_ratio, student_step)\n exp_logger.add_scalar(\"student/grad_norm\", student_grad_stats.get(\"grad_norm\", 0.0), student_step)\n for k, v in reduced_log_info.items():\n exp_logger.add_scalar(k, v, student_step)\n exp_logger.add_scalar(\"global/peak_gpu_memory_gb\", torch.cuda.max_memory_allocated(device) / (1024 ** 3), student_step)\n\n student_step += 1\n info_accumulated = defaultdict(list)\n\n if config.train.student.ema and ema is not None and accelerator.sync_gradients:\n # FSDP2: EMA operates on local shards via _get_data() in EMAModuleWrapper\n ema.step(transformer_trainable_parameters, student_step)\n\n # End of epoch: cleanup to prevent memory fragmentation accumulation\n # torch.cuda.empty_cache()\n\n # Final checkpoint save (skip if already saved by the periodic checkpoint logic)\n final_checkpoint_path = os.path.join(checkpoint_dir, f\"checkpoint-{student_step}\")\n already_saved = os.path.exists(final_checkpoint_path)\n if not already_saved:\n if accelerator.is_main_process:\n os.makedirs(final_checkpoint_path, exist_ok=True)\n accelerator.wait_for_everyone()\n accelerator.save_state(final_checkpoint_path)\n if ema is not None and accelerator.is_main_process:\n torch.save(ema.state_dict(), os.path.join(final_checkpoint_path, \"ema_state.pt\"))\n # FSDP: save full (unsharded) state dicts for prepare_student_pipeline compatibility\n save_fsdp_full_checkpoint(accelerator, transformer, ema, final_checkpoint_path, logger)\n if accelerator.is_main_process:\n logger.info(f\"Saved final checkpoint to {final_checkpoint_path}\")\n else:\n if accelerator.is_main_process:\n logger.info(f\"Final checkpoint already exists at {final_checkpoint_path}, skipping duplicate save\")\n\n if accelerator.is_main_process:\n logger.info(f\"Training completed in {time.time() - train_start_time:.2f} seconds\")\n # Write marker file so submit.py can locate the experiment directory\n marker_path = os.path.join(config.output_dir, f\".experiment_dir\")\n with open(marker_path, \"w\") as f:\n f.write(experiment_dir)\n logger.info(f\"Wrote experiment marker to {marker_path}\")\n\n accelerator.wait_for_everyone()\n if torch.distributed.is_initialized():\n torch.distributed.destroy_process_group()\n shutdown_image_save_executor()\n\nif __name__ == \"__main__\":\n app.run(main)"}}}} |