code/src/model_training/training_modules/training_loop.py.nanguard-bak

import logging
import math
import os
from typing import Optional

import torch
import torch.distributed as dist
from accelerate import Accelerator
from tqdm import tqdm

from src.model_training.transformers_compat import patch_transformers_hybrid_cache

patch_transformers_hybrid_cache()

from diffsynth.trainers.utils import DiffusionTrainingModule
from src.model_training.fov_retrieval import FOVMemoryRetriever
from src.model_training.fov_retrieval import retrieve_context_frames_advanced, retrieve_fov_context_frames
from src.model_training.training_modules.model_logger import ModelLogger

logger = logging.getLogger(__name__)


def launch_training_task(
    dataset: torch.utils.data.Dataset,
    model: DiffusionTrainingModule,
    model_logger: ModelLogger,
    optimizer: torch.optim.Optimizer,
    scheduler: torch.optim.lr_scheduler.LRScheduler,
    num_epochs: int = 1,
    gradient_accumulation_steps: int = 1,
    per_device_train_batch_size: int = 1,
    seed: int = 42,
    spike_threshold: float = 5.0,
    resume_step_count: int = 0,
    enable_fov_retrieval: bool = False,
    retrieval_method: str = "fov",  # fov | latent_sim
    latent_retrieval_dir: Optional[str] = None,
    dataset_base_path: str = None,
    fov_retriever: Optional[FOVMemoryRetriever] = None,
    context_memory_frames: int = 5,
    prev_chunk_frames: int = 81,
    fov_top_k: int = 4,  # Number of overlap frames to retrieve. GT frame 0 will be added automatically.
    use_rt_relative: bool = False,  # Experiment 1_4_2: Use RT relative conversion (aligned with Context-as-Memory)
    strict_overlap_context: bool = False,
    dataset_repeat: int = 1,  # Add dataset_repeat parameter for step calculation
    use_camera_encoder: bool = False,  # exp1_4_3: use CameraEncoder (action_mlp unused -> need find_unused_parameters)
    num_workers: int = 0,  # DataLoader workers: 0=main process, >0=parallel preload (recommend 4 for video)
    context_source: str = "fov",
    max_train_steps: int = 0,
    progress_total_steps: int = 0,
    log_interval: int = 20,
    max_grad_norm: float = 1.0,
):
    prev_chunk_frames = int(prev_chunk_frames)
    # VideoDataset can return None when file loading fails; keep distributed batches aligned.
    def collate_fn(batch):
        valid_batch = [item for item in batch if item is not None]
        return valid_batch or None
    
    num_workers = max(0, int(num_workers))
    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=per_device_train_batch_size,
        shuffle=True,
        collate_fn=collate_fn,
        num_workers=num_workers,
        drop_last=True,
        persistent_workers=(num_workers > 0),
        pin_memory=(num_workers > 0 and torch.cuda.is_available()),
    )
    if num_workers > 0:
        logger.info(f"[DataLoader] num_workers={num_workers}, persistent_workers=True, pin_memory={torch.cuda.is_available()} (data preload parallel to GPU)")
    
    timeout_seconds = int(os.environ.get('TORCH_DISTRIBUTED_DEFAULT_TIMEOUT', 2400))
    os.environ['TORCH_DISTRIBUTED_DEFAULT_TIMEOUT'] = str(timeout_seconds)
    logger.info(f"[Timeout Config] Setting TORCH_DISTRIBUTED_DEFAULT_TIMEOUT={timeout_seconds} seconds ({timeout_seconds/60:.1f} minutes)")
    
    # Conditional context paths can leave parameters unused on some iterations.
    need_find_unused = bool(use_camera_encoder) or model_logger.context_drop_prob > 0.0
    if need_find_unused:
        from accelerate import DistributedDataParallelKwargs
        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
        accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps, kwargs_handlers=[ddp_kwargs])
        logger.info("[DDP] find_unused_parameters=True (conditional modules / context_drop_prob enabled)")
    else:
        accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps)
    model, optimizer, dataloader, scheduler = accelerator.prepare(model, optimizer, dataloader, scheduler)
    
    if model_logger.enable_video_sampling and model_logger.total_steps is not None:
        dataset_size = len(dataset)
        num_processes = accelerator.num_processes
        effective_dataset_size = dataset_size * dataset_repeat
        total_steps_per_gpu = (effective_dataset_size * num_epochs) // (gradient_accumulation_steps * num_processes * per_device_train_batch_size)
        total_steps_global = total_steps_per_gpu * num_processes
        model_logger.total_steps = total_steps_global
        
        if accelerator.is_main_process:
            logger.info("="*80)
            logger.info("[Step Calculation] Corrected total_steps after accelerator.init")
            logger.info("="*80)
            logger.info(f"  Dataset size (unique samples): {dataset_size}")
            logger.info(f"  Dataset repeat: {dataset_repeat}")
            logger.info(f"  Effective dataset size: {effective_dataset_size} (unique * repeat)")
            logger.info(f"  Number of epochs: {num_epochs}")
            logger.info(f"  Number of GPUs: {num_processes}")
            logger.info(f"  Gradient accumulation steps: {gradient_accumulation_steps}")
            logger.info(f"  Per-device batch size: {per_device_train_batch_size}")
            logger.info(f"  Total samples to process: {effective_dataset_size * num_epochs}")
            logger.info(f"  Steps per GPU: ~{total_steps_per_gpu}")
            logger.info(f"  Total steps (global): {total_steps_global}")
            logger.info("")
            logger.info(f"  ✓ Each GPU will process ~{total_steps_per_gpu} steps")
            logger.info(f"  ✓ This ensures traversal of all {effective_dataset_size} samples")
            logger.info(f"    ({dataset_size} unique samples × {dataset_repeat} repeats)")
            logger.info(f"  ✓ Over {num_epochs} epoch(s)")
            logger.info("="*80)
    
    step = resume_step_count
    traj_loss = 0.0
    if resume_step_count > 0:
        adaptation_steps = max(200, resume_step_count // 100)
        spike_detection_start_step = resume_step_count + adaptation_steps
        logger.info(f"Resuming from step {resume_step_count}, spike detection will start at step {spike_detection_start_step} (after {adaptation_steps} adaptation steps)")
    else:
        spike_detection_start_step = 100

    for epoch_id in range(num_epochs):
        epoch_seed = seed + epoch_id
        torch.manual_seed(epoch_seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed(epoch_seed)
            torch.cuda.manual_seed_all(epoch_seed)
        
        if resume_step_count > 0 and epoch_id == 0:
            estimated_skip = resume_step_count // gradient_accumulation_steps
            if estimated_skip > 0:
                logger.info(f"Skipping {estimated_skip} data samples to resume from step {resume_step_count}...")
                dataloader_iter = iter(dataloader)
                for _ in tqdm(range(estimated_skip), desc="Skipping data", unit="samples", leave=False):
                    try:
                        next(dataloader_iter)
                    except StopIteration:
                        break
                dataloader = dataloader_iter
                logger.info(f"Successfully skipped {estimated_skip} data samples, resuming training...")
        
        # Track consecutive None data to detect if we're stuck in a loop
        consecutive_none_count = 0
        max_consecutive_none = 100  # If we get 100 consecutive None values, something is wrong
        
        progress_total = int(progress_total_steps)
        if progress_total <= 0:
            progress_total = len(dataloader)
        progress_bar = tqdm(
            dataloader,
            total=progress_total,
            initial=resume_step_count if progress_total_steps else 0,
            desc="Training steps",
            unit="step",
        )
        for data_idx, data in enumerate(progress_bar):
            # Handle None data (can happen if all files in batch fail to load)
            if data is None:
                consecutive_none_count += 1
                if consecutive_none_count >= max_consecutive_none:
                    logger.error(f"Received {max_consecutive_none} consecutive None data samples. This suggests a serious dataset issue. Stopping training.")
                    raise ValueError(f"Too many consecutive None data samples ({max_consecutive_none}). Check dataset files.")
                
                # Log warning but continue (will skip this step)
                if consecutive_none_count <= 10 or consecutive_none_count % 10 == 0:
                    logger.warning(f"Received None data at index {data_idx} (consecutive: {consecutive_none_count}). This may indicate missing or corrupted files. Skipping...")
                
                # Still increment step to keep step_count synchronized
                step += 1
                dummy_loss = torch.tensor(0.0, device=accelerator.device, requires_grad=False)
                model_logger.on_step_end(dummy_loss, accelerator, model, current_batch=samples)
                continue
            
            # Reset consecutive None counter when we get valid data
            consecutive_none_count = 0
            
            # Normalize to list of samples for batch processing (per_device_train_batch_size > 1)
            samples = data if isinstance(data, list) else [data]
            
            # Simplified context-based retrieval OR replay/prev_chunk_tail (aligned with multichunk eval)
            context_retrieval_success = True  # Set False if any sample fails (for strict mode)
            _umodel = accelerator.unwrap_model(model)
            _cm_frames = int(_umodel.context_memory_frames)
            _cs = context_source.strip().lower()
            if _cs not in ("fov", "replay", "prev_chunk_tail"):
                _cs = "fov"

            if _cs == "replay" and dataset_base_path:
                from src.model_training.multichunk_sample_utils import (
                    replay_context_actions_from_segment_actions,
                    replay_context_global_indices,
                    synthetic_replay_context_from_segment,
                )
                for d in samples:
                    vf = d.get("video") or []
                    n_seg = min(int(prev_chunk_frames), len(vf)) if vf else 0
                    ctx_pil = synthetic_replay_context_from_segment(vf, n_seg, _cm_frames) if n_seg > 0 else None
                    if not ctx_pil:
                        context_retrieval_success = False
                        break
                    d["context_frames"] = ctx_pil
                    d["context_source"] = "replay_synthetic"
                    acts = d.get("actions")
                    if isinstance(acts, list) and len(acts) >= n_seg:
                        ra = replay_context_actions_from_segment_actions(acts[:n_seg], n_seg, _cm_frames)
                        if ra is not None:
                            d["context_actions"] = ra
                    sf = int(d.get("start_frame", 0) or 0)
                    idxs = replay_context_global_indices(n_seg, _cm_frames)
                    d["context_frame_indices"] = [sf + int(i) for i in idxs]

            elif _cs == "prev_chunk_tail" and dataset_base_path:
                from src.model_training.multichunk_sample_utils import load_prev_chunk_tail_from_disk, load_prev_chunk_tail_rt_actions
                _ctx_pos = os.environ.get("CONTEXT_POSITION", "suffix").strip().lower()
                _nearest_first = (_ctx_pos == "suffix")
                for d in samples:
                    sf = int(d.get("start_frame", 0) or 0)
                    vn = d.get("video_name", "")
                    pil_list, idxs = load_prev_chunk_tail_from_disk(
                        dataset_base_path, str(vn), sf, _cm_frames, nearest_first=_nearest_first
                    )
                    if not pil_list:
                        context_retrieval_success = False
                        break
                    d["context_frames"] = pil_list
                    d["context_frame_indices"] = list(idxs) if idxs else []
                    d["context_source"] = "prev_chunk_tail"
                    ra, _ = load_prev_chunk_tail_rt_actions(
                        dataset_base_path,
                        str(vn),
                        sf,
                        _cm_frames,
                        use_rt_relative=use_rt_relative,
                        nearest_first=_nearest_first,
                    )
                    if ra:
                        d["context_actions"] = ra

            elif enable_fov_retrieval and dataset_base_path:
                for d in samples:
                    if retrieval_method == "latent_sim":
                        (
                            context_frames,
                            context_actions,
                            context_indices,
                            ref_frame_idx,
                            video_name,
                            source,
                        ) = retrieve_context_frames_advanced(
                            data=d,
                            dataset_base_path=dataset_base_path,
                            top_k=fov_top_k,
                            drop_overlap_probability=0.1,
                            use_rt_relative=use_rt_relative,
                            retrieval_method="latent_sim",
                            latent_retrieval_dir=latent_retrieval_dir,
                            strict_overlap_labels=strict_overlap_context,
                        )
                    else:
                        (
                            context_frames,
                            context_actions,
                            context_indices,
                            ref_frame_idx,
                            video_name,
                            source,
                        ) = retrieve_fov_context_frames(
                            data=d,
                            dataset_base_path=dataset_base_path,
                            fov_retriever=fov_retriever,  # unused in simplified retrieval, kept for compat
                            top_k=fov_top_k,  # fov_top_k is number of overlap frames (4), GT frame 0 will be added automatically
                            use_precomputed_overlaps=True,
                            strict_overlap_labels=strict_overlap_context,
                            allow_realtime_fallback=(not strict_overlap_context),
                            allow_segment_fallback=(not strict_overlap_context),
                        )

                    if context_frames and len(context_frames) > 0:
                        # Use retrieved frames as context
                        d["context_frames"] = context_frames
                        if context_actions:
                            d["context_actions"] = context_actions
                        # Store retrieval metadata for visualization/debugging
                        d["context_frame_indices"] = context_indices
                        d["context_ref_frame_idx"] = ref_frame_idx
                        d["context_video_name"] = video_name
                        d["context_source"] = source
                    else:
                        context_retrieval_success = False
                        break
                        
            # Strict mode: if we require context but retrieval failed, skip this step
            _need_ctx_strict = (
                strict_overlap_context
                and (not context_retrieval_success)
                and (
                    enable_fov_retrieval
                    or context_source.strip().lower() in ("replay", "prev_chunk_tail")
                )
            )
            if _need_ctx_strict:
                if step % 50 == 0 and accelerator.is_main_process:
                    logger.warning(f"[CONTEXT][STRICT] No context at step={step}, skipping this training sample.")
                step += 1
                dummy_loss = torch.tensor(0.0, device=accelerator.device, requires_grad=False)
                model_logger.on_step_end(dummy_loss, accelerator, model, current_batch=samples)
                continue

            # Diagnostic: cache the raw dataset video + a VAE encode→decode roundtrip
            # every CACHE_INPUT_VIDEO_EVERY steps (env, default 0=off) to check whether
            # the videos the model trains on have color issues (inversion / R-B swap).
            # Main-process only; no_grad side-effect — never affects the forward/grads.
            _cache_every = int(os.environ.get("CACHE_INPUT_VIDEO_EVERY", "0") or "0")
            if _cache_every > 0 and accelerator.is_main_process and (step % _cache_every == 0):
                try:
                    _dbg_dir = os.path.join(model_logger.output_path, "input_video_debug")
                    _umodel.dump_input_video_debug(samples, step, _dbg_dir)
                except Exception as _e:  # noqa: BLE001 - diagnostic only
                    print(f"[input-video-debug] step={step} skipped: {_e}", flush=True)

            with accelerator.accumulate(model):
                optimizer.zero_grad()
                # One forward over full batch: data is list of B dicts when per_device_train_batch_size > 1
                # Main loss on current batch
                loss = model(data)

                step += 1
                if traj_loss == 0.0:
                    traj_loss = loss.item()
                else:
                    alpha = 0.01
                    traj_loss = (1 - alpha) * traj_loss + alpha * loss.item()
                
                if step >= spike_detection_start_step and traj_loss > 0:
                    relative_loss = loss.item() / traj_loss
                    if resume_step_count > 0 and step < resume_step_count + 500:
                        effective_threshold = spike_threshold * 1.5
                    else:
                        effective_threshold = spike_threshold
                    
                    should_skip = relative_loss > effective_threshold
                    # Keep the skip decision identical across ranks to avoid DDP hangs.
                    skip_t = torch.tensor(1.0 if should_skip else 0.0, device=accelerator.device, dtype=torch.float32)
                    if accelerator.num_processes > 1:
                        dist.all_reduce(skip_t, op=dist.ReduceOp.MAX)
                    skip_global = skip_t.item() > 0.5
                    
                    if skip_global:
                        if accelerator.is_main_process:
                            logger.warning(f"Spike detected at step {step} (loss={loss.item():.4f}, traj_loss={traj_loss:.4f}, ratio={relative_loss:.2f}), sync skip across all ranks")
                        dummy_loss = torch.tensor(0.0, device=accelerator.device, requires_grad=False)
                        model_logger.on_step_end(dummy_loss, accelerator, model, current_batch=samples)
                        del loss
                        torch.cuda.empty_cache()
                        continue
                accelerator.backward(loss)
                loss_v = float(loss.detach().float().item())
                # CGLA SSE load-balancing aux loss (only when
                # cgla_aux_loss_weight > 0). Each CGLATransformerBlock stashes
                # its aux_loss during the forward; sum across blocks (mirrors
                # WanVideoPipeline.training_loss). Reported for monitoring — it
                # is already added to `loss` inside the pipeline.
                aux_v = float("nan")
                _aux_w = 0.0
                _umodel = accelerator.unwrap_model(model)
                _pipe = getattr(_umodel, "pipe", None)
                if _pipe is not None:
                    _aux_w = float(getattr(_pipe, "cgla_aux_loss_weight", 0.0) or 0.0)
                if _aux_w > 0:
                    _aux_sum = None
                    for _b in getattr(getattr(_pipe, "dit", None), "blocks", []):
                        _a = getattr(_b, "_aux_loss", None)
                        if _a is not None and torch.is_tensor(_a) and torch.isfinite(_a).all():
                            _aux_sum = _a if _aux_sum is None else _aux_sum + _a
                    if _aux_sum is not None:
                        aux_v = float(_aux_sum.detach().float().item())
                # grad_norm / clip only on the accumulation boundary (sync step):
                # avoid a wasteful all-reduce every micro-step, and measure the
                # full accumulated gradient rather than a partial one.
                grad_norm = float("nan")
                cur_lr = float(optimizer.param_groups[0].get("lr", 0.0))
                if accelerator.sync_gradients:
                    if max_grad_norm and max_grad_norm > 0:
                        grad_norm_t = accelerator.clip_grad_norm_(model.parameters(), max_grad_norm)
                    else:
                        grad_norm_t = accelerator.clip_grad_norm_(model.parameters(), float("inf"))
                    grad_norm = float(grad_norm_t.detach().float().item())
                    cur_lr = float(optimizer.param_groups[0].get("lr", 0.0))
                optimizer.step()
                # Live progress bar: loss every step; grad_norm/lr on the sync step;
                # aux when present.
                _post = {"loss": f"{loss_v:.4f}", "traj": f"{traj_loss:.4f}"}
                if not math.isnan(grad_norm):
                    _post["gnorm"] = f"{grad_norm:.3f}"
                    _post["lr"] = f"{cur_lr:.2e}"
                if not math.isnan(aux_v):
                    _post["aux"] = f"{aux_v:.4f}"
                progress_bar.set_postfix(_post)
                # Periodic structured line (main process, sync step only) so
                # loss/grad_norm/aux land in the tee'd log file, not just the bar.
                if (accelerator.sync_gradients and accelerator.is_main_process
                        and (step % log_interval == 0)):
                    _aux_str = (f" aux={aux_v:.4f}(w={_aux_w:.3f})"
                                if not math.isnan(aux_v) else "")
                    logger.info(
                        f"[train] step={step} loss={loss_v:.4f} "
                        f"traj={traj_loss:.4f} grad_norm={grad_norm:.4f} "
                        f"lr={cur_lr:.2e}{_aux_str}"
                    )
                model_logger.on_step_end(loss, accelerator, model, current_batch=samples,
                                         aux_loss=aux_v)
                scheduler.step()

                if max_train_steps and step >= max_train_steps:
                    if progress_total_steps:
                        progress_bar.n = min(step, progress_bar.total) if progress_bar.total is not None else step
                        progress_bar.refresh()
                    if accelerator.is_main_process:
                        logger.info(f"[TRAIN] Reached max_train_steps={max_train_steps}; stopping without epoch checkpoint.")
                    accelerator.wait_for_everyone()
                    return
            
        model_logger.on_epoch_end(accelerator, model, epoch_id)