{"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)"}}}}