train_matdiff.py

import argparse
import logging
import os
from pathlib import Path

os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import torch
import torch.utils.checkpoint

from einops import rearrange
import transformers
from accelerate import DistributedType
from accelerate.logging import get_logger
from accelerate.utils import set_seed
from tqdm.auto import tqdm

import diffusers
from diffusers import (
    FlowMatchEulerDiscreteScheduler,
    SD3Transformer2DModel,
)
from diffusers.optimization import get_scheduler, get_cosine_schedule_with_warmup
from diffusers.training_utils import (
    cast_training_params,
    compute_loss_weighting_for_sd3,
)
from mcgen.utils.sharded_ema import ShardedEMA

from mcgen.matdiff.training.logging import (
    log_training_batch,
    log_latent_verification,
)
from mcgen.training.image_logging import channel_strip_images, log_tracker_images_and_scalars

from mcgen.training.checkpointing import (
    CheckpointPolicy,
    checkpoint_step,
    copy_state_to_named_parameters,
)
from mcgen.matdiff.training.batch import prepare_matdiff_training_batch, prepare_matdiff_validation_batch
from mcgen.matdiff.training.data import build_matdiff_train_dataloader, build_validation_data
from mcgen.matdiff.training.flow import (
    apply_flow_preconditioning,
    flow_matching_target,
    lookup_sigmas,
    sample_timestep_density,
    weighted_flow_mse,
)
from mcgen.matdiff.training.optimizers import build_matdiff_optimizer
from mcgen.matdiff.training.state import (
    register_matdiff_checkpoint_hooks,
    save_matdiff_ema_state,
    save_matdiff_transformer_state,
    unwrap_compiled_module,
)
from mcgen.training.runtime import StepPerformanceTracker, build_wandb_accelerator, training_dtype_config
from mcgen.utils.print_utils import print_info, print_info_main
from mcgen.matdiff.conditioning import append_condition_channels
from mcgen.matdiff.latents import generate_latents, decode_latents_to_images
from mcgen.matdiff.mode import MatDiffModelMode
from mcgen.matdiff.weights import initialize_transformer_weights
from mcgen.matvae.loading import load_offset_autoencoder
from mcgen.matdiff.attention_payload import payload_to_device
from mcgen.matdiff.attention_processors import collect_attention_sites
from mcgen.matdiff.attention_correspondence import attach_corr_to_batch
from mcgen.matdiff.attention_processor_lora import LoRAMessagePassingAttnProcessor
from mcgen.matdiff.model_loader import apply_lora_processors, extend_pos_embed_in, build_peft_lora_config
from mcgen.utils.correspondence import correspondence_psnr
from mcgen.utils.config import load_config
from omegaconf import OmegaConf

logger = get_logger(__name__)


def _drop_precomputed_encoder_modules(module) -> None:
    for attr in ("encoder", "quant_conv"):
        if hasattr(module, attr):
            delattr(module, attr)


def _prepare_autoencoders_for_training(
    vae,
    matvae,
    *,
    use_dual_branch: bool,
    device: torch.device,
    dtype: torch.dtype,
) -> None:
    """Keep only decode-time VAE components on devices needed by MatDiff training."""
    vae.to(device, dtype=dtype)
    vae.decoder.to("cpu")
    _drop_precomputed_encoder_modules(vae)

    if use_dual_branch:
        matvae.to(torch.device("cpu"))
        return

    matvae.to(device, dtype=dtype)
    matvae.decoder.to("cpu")
    _drop_precomputed_encoder_modules(matvae)


def _set_validation_decoders_device(vae, matvae, *, use_dual_branch: bool, device) -> None:
    vae.decoder.to(device)
    if not use_dual_branch:
        matvae.decoder.to(device)


def matdiff_validation_enabled(train_cfg) -> bool:
    validation_steps = getattr(train_cfg, "validation_steps", None)
    return validation_steps is not None and int(validation_steps) > 0


def matdiff_final_validation_enabled(train_cfg, *, validation_enabled: bool) -> bool:
    return bool(getattr(train_cfg, "run_final_validation", validation_enabled)) and validation_enabled


def matdiff_is_full_finetune(transformer_lora_config, *, attn_lora_enabled: bool) -> bool:
    """Use full transformer fine-tuning whenever PEFT LoRA is not configured.

    ``model.attn_lora`` controls MatDiff's custom multi-view attention
    processors, not the PEFT adapter training mode.  The legacy MCGen training
    path fine-tuned the full transformer for ``lora.rank == 0`` even when these
    attention processors were enabled.
    """
    del attn_lora_enabled
    return transformer_lora_config is None


def main(cfg):
    logging_dir = Path(cfg.project.output_dir, cfg.project.logging_dir)

    accelerator = build_wandb_accelerator(
        project_dir=cfg.project.output_dir,
        logging_dir=logging_dir,
        gradient_accumulation_steps=cfg.train.gradient_accumulation_steps,
        mixed_precision=cfg.train.mixed_precision,
        find_unused_parameters=False,
    )

    # Make one log on every process with the configuration for debugging.
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
    )
    logger.info(accelerator.state, main_process_only=False)
    if accelerator.is_local_main_process:
        transformers.utils.logging.set_verbosity_info()
        diffusers.utils.logging.set_verbosity_info()
    else:
        transformers.utils.logging.set_verbosity_error()
        diffusers.utils.logging.set_verbosity_error()

    # If passed along, set the training seed now.
    if cfg.seed is not None:
        set_seed(cfg.seed)

    # Create output directory
    if accelerator.is_main_process:
        os.makedirs(cfg.project.output_dir, exist_ok=True)

    # Load scheduler and models (text encoders removed – embeddings precomputed in dataset)
    noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
        cfg.model.pretrained_model_name_or_path, subfolder="scheduler"
    )
    # Load base RGB VAE + MatVAE
    matvae_offset_mode = getattr(cfg.model, "matvae_offset_mode", 'scaled')
    vae, matvae = load_offset_autoencoder(
        cfg.model.sd_vae_path,
        cfg.model.matvae_config,
        cfg.model.matvae_ckpt,
        offset_mode=matvae_offset_mode,
    )
    transformer = SD3Transformer2DModel.from_pretrained(
        cfg.model.pretrained_model_name_or_path, subfolder="transformer"
    )
    vae.eval().requires_grad_(False)
    matvae.eval().requires_grad_(False)

    transformer.requires_grad_(False)
    vae.requires_grad_(False)
    matvae.requires_grad_(False)

    mode = MatDiffModelMode.from_config(cfg)
    num_views = mode.num_views
    use_dual_branch = mode.use_dual_branch
    use_caa = mode.use_caa
    use_rope = mode.use_rope
    use_weaving_attention = mode.use_weaving_attention
    corr_dilate_iterations = mode.corr_dilate_iterations
    gradual_nheads = mode.gradual_nheads(transformer.config.num_attention_heads)
    condition_channels = mode.condition_channels

    # ------------------------------------------------------------------
    # Replace attention processors with LoRAMessagePassingAttnProcessor (using shared function)
    apply_lora_processors(
        transformer,
        cfg.model,
        use_caa=use_caa,
        use_rope=use_rope,
        num_views=num_views,
        num_domains=mode.num_domains,
        use_global_pos=mode.use_global_pos,
    )
    attn_lora_enabled = bool(
        getattr(cfg.model, "attn_lora", None) and getattr(cfg.model.attn_lora, "enabled", False)
    )
    if attn_lora_enabled:
        logger.info(
            "Activated LoRAMessagePassingAttnProcessor via apply_lora_processors()"
        )


    # ╭──────────────────────────────── DTYPE POLICY ───────────────────────────────╮
    # │ See training_dtype_config() in mcgen.training.runtime for details.          │
    # ╰─────────────────────────────────────────────────────────────────────────────╯
    transformer_lora_config = build_peft_lora_config(cfg.lora)
    is_full_finetune = matdiff_is_full_finetune(transformer_lora_config, attn_lora_enabled=attn_lora_enabled)
    dtype_cfg = training_dtype_config(cfg.train.mixed_precision, is_full_finetune=is_full_finetune)
    weight_dtype = dtype_cfg.weight_dtype
    keep_trainable_fp32 = dtype_cfg.keep_trainable_fp32

    _prepare_autoencoders_for_training(
        vae,
        matvae,
        use_dual_branch=use_dual_branch,
        device=accelerator.device,
        dtype=weight_dtype,
    )
    torch.cuda.empty_cache()
    transformer.to(accelerator.device, dtype=weight_dtype)

    # --- Optionally expand transformer input channels for conditioning ---
    if mode.has_conditioning:
        old_in = transformer.pos_embed.proj.in_channels
        extend_pos_embed_in(transformer, condition_channels)
        new_in = transformer.pos_embed.proj.in_channels
        logger.info(f"Expanded transformer.pos_embed.proj in_channels {old_in} -> {new_in} (added {condition_channels} cond channels; zero initialized)")


    if cfg.train.gradient_checkpointing:
        transformer.enable_gradient_checkpointing()
    
    if is_full_finetune:
        transformer.requires_grad_(True)
    else:
        if transformer_lora_config is None:
            raise RuntimeError("Non-full MatDiff training requires a PEFT LoRA config.")

        target_modules = list(transformer_lora_config.target_modules)
        blocks_str = cfg.lora.lora_blocks if cfg.lora.lora_blocks is not None else "all"
        sample_targets = ", ".join(target_modules[:4]) + (" ..." if len(target_modules) > 4 else "")
        print_info_main(
            f"LoRA config -> rank: {cfg.lora.rank}, alpha: {cfg.lora.rank}, dropout: {cfg.lora.lora_dropout}, "
            f"blocks: {blocks_str}, target modules: {len(target_modules)} (e.g., {sample_targets})"
        )
        transformer.add_adapter(transformer_lora_config)

        if mode.has_conditioning:
            for p in transformer.pos_embed.proj.parameters():
                p.requires_grad = True
            print_info_main(
                "Enabled gradients for transformer.pos_embed.proj to learn conditioning channels under LoRA fine-tuning."
            )

        # make attention processor params trainable
        for _, attn_mod, _ in collect_attention_sites(transformer):
            proc = attn_mod.get_processor()
            if isinstance(proc, torch.nn.Module):
                for p in proc.parameters():
                    p.requires_grad = True

    @torch.no_grad()
    def _run_validation_step(current_epoch_or_step: int, val_data, *, is_final: bool = False):
        was_training = transformer.training
        if was_training:
            transformer.eval()
        
        _set_validation_decoders_device(
            vae,
            matvae,
            use_dual_branch=use_dual_branch,
            device=accelerator.device,
        )
        try:
            validation_batch = prepare_matdiff_validation_batch(
                val_data,
                num_views=cfg.model.num_views,
                device=accelerator.device,
                dtype=weight_dtype,
                use_dual_branch=use_dual_branch,
                null_prompt_embeds=null_prompt_embeds,
                null_pooled_prompt_embeds=null_pooled_prompt_embeds,
            )
            layout = validation_batch.geometry

            attn_payload = None
            if use_caa:
                attn_payload = payload_to_device(
                    val_data['attn_payload'], accelerator.device, pos_dtype=weight_dtype)

            latents = generate_latents(
                transformer=unwrap_compiled_module(transformer),
                noise_scheduler=noise_scheduler,
                prompt_embeds=validation_batch.prompt_embeds,
                pooled_prompt_embeds=validation_batch.pooled_prompt_embeds,
                negative_prompt_embeds=validation_batch.negative_prompt_embeds,
                negative_pooled_prompt_embeds=validation_batch.negative_pooled_prompt_embeds,
                height=cfg.model.resolution,
                width=cfg.model.resolution,
                num_inference_steps=20,
                guidance_scale=4.0,
                weight_dtype=weight_dtype,
                device=accelerator.device,
                condition_channels=condition_channels,
                condition_latents=validation_batch.cond_latents,
                attn_payload=attn_payload,
            )

            # Decode latents to 5-channel images using shared function
            # Pass raw cond_values for control signal (only used when use_dual_branch=False)
            cond_for_decode = validation_batch.cond_values if not use_dual_branch else None
            pred_combined_flat = decode_latents_to_images(
                latents,
                vae,
                matvae,
                vae_scaling_factor=vae_config_scaling_factor,
                vae_shift_factor=vae_config_shift_factor,
                use_dual_branch=use_dual_branch,
                num_views=layout.views,
                num_groups=layout.groups,
                cond_values=cond_for_decode,
            )

            pred_combined_flat = pred_combined_flat.to(device=accelerator.device, dtype=weight_dtype)
            pred_rgb = pred_combined_flat[:, :3]
            pred_mr = pred_combined_flat[:, 3:]

            corr_ell_full = val_data['corr_ell_full']  # (bVHW, K)
            rgb_psnr, _, _ = correspondence_psnr(pred_rgb, corr_ell_full, data_range=2.0)
            mr_psnr, _, _ = correspondence_psnr(pred_mr, corr_ell_full, data_range=2.0)

            gt_rgb = val_data['pixel_values']
            if use_dual_branch:
                gt_rgb = gt_rgb[0::2]
            gt_rgb = gt_rgb[:, :, :3]
            gt_rgb = rearrange(gt_rgb, 'b V C H W -> (b V) C H W').to(device=accelerator.device, dtype=weight_dtype)
            gt_rgb_psnr, _, _ = correspondence_psnr(gt_rgb, corr_ell_full, data_range=2.0)

            vis_tensor = pred_combined_flat
            if mode.has_conditioning:
                cond_for_decode = validation_batch.cond_latents  # (bD*V, C, h, w)
                if use_dual_branch:
                    cond_for_decode = rearrange(
                        cond_for_decode,
                        '(b D V) ... -> b D V ...',
                        b=layout.groups,
                        D=2,
                        V=layout.views,
                    )[:, 0]
                    cond_for_decode = rearrange(cond_for_decode, 'b V ... -> (b V) ...')
                inv_cond = (cond_for_decode / vae_config_scaling_factor) + vae_config_shift_factor
                inv_cond = inv_cond.to(device=vae.device, dtype=vae.dtype)
                latent_channels = vae.config.latent_channels
                if inv_cond.shape[1] % latent_channels != 0:
                    raise ValueError(
                        f"condition_channels {inv_cond.shape[1]} not multiple of VAE latent channels "
                        f"{latent_channels}"
                    )
                dec_cond = []
                with torch.no_grad():
                    for i in range(0, inv_cond.shape[1], latent_channels):
                        slice_cond = inv_cond[:, i:i+latent_channels]
                        dec_cond.append(vae.decode(slice_cond).sample)
                dec_cond = torch.cat(dec_cond, dim=1).to(device=accelerator.device, dtype=weight_dtype)
                vis_tensor = torch.cat([dec_cond, vis_tensor], dim=1)

            grids = channel_strip_images(vis_tensor.detach().cpu())
            phase = "test" if is_final else "validation"
            log_tracker_images_and_scalars(
                accelerator,
                step=current_epoch_or_step,
                images={phase + "_grid": grids},
                scalars={
                    phase + "_psnr_rgb": float(rgb_psnr),
                    phase + "_psnr_mr": float(mr_psnr),
                    phase + "_psnr_rgb_gt": float(gt_rgb_psnr),
                },
            )
        finally:
            _set_validation_decoders_device(
                vae,
                matvae,
                use_dual_branch=use_dual_branch,
                device="cpu",
            )

            if was_training:
                transformer.train()

    # load partial transformer weights if provided
    initialize_transformer_weights(transformer, cfg.model.init_transformer_weights, logger)

    # Enable TF32 for faster training on Ampere GPUs,
    if cfg.train.allow_tf32 and torch.cuda.is_available():
        try:
            torch.backends.cuda.matmul.fp32_precision = "tf32"
            print_info("Enabled TF32 precision for matmul and convolutions.")
        except Exception:
            print_info("Failed to enable TF32 precision; trying with legacy api.")
            torch.set_float32_matmul_precision('high')

    if cfg.optim.scale_lr:
        # Effective batch experienced by model multiplies by num_views (flattened)
        cfg.optim.learning_rate = (
            cfg.optim.learning_rate
            * cfg.train.gradient_accumulation_steps
            * cfg.data.train.batch_size
            * accelerator.num_processes
            * cfg.model.num_views
        )

    proc_modules = []
    for m in transformer.modules():
        if isinstance(m, LoRAMessagePassingAttnProcessor):
            proc_modules.append(m)

    # Make sure the trainable lora params are in float32
    if keep_trainable_fp32:
        cast_training_params([transformer], dtype=torch.float32)

    # Make sure the attention processor params are in float32
    if proc_modules:
        cast_training_params(proc_modules, dtype=torch.float32)

    transformer_lora_parameters = [p for p in transformer.parameters() if p.requires_grad]
    use_ema = bool(getattr(cfg.train, 'use_ema', False))
    ema_model = None
    optimizer = build_matdiff_optimizer(cfg.optim, transformer_lora_parameters, logger)

    # Dataset and DataLoaders creation:
    print_info_main("Loading rgbmr_mv dataset with precomputed embeddings ...")

    validation_enabled = matdiff_validation_enabled(cfg.train)
    val_data = build_validation_data(
        cfg,
        accelerator=accelerator,
        enabled=validation_enabled,
        use_dual_branch=use_dual_branch,
        use_global_pos=mode.use_global_pos,
        use_weaving_attention=use_weaving_attention,
        corr_dilate_iterations=corr_dilate_iterations,
        gradual_nheads=gradual_nheads,
    )
    train_dataloader = build_matdiff_train_dataloader(cfg, use_dual_branch=use_dual_branch)

    # load null embeddings (required; shapes already (1,T,D)/(1,Dp))
    null_prompt_embeds = torch.load(cfg.project.null_pe).clone()
    null_pooled_prompt_embeds = torch.load(cfg.project.null_ppe).clone()
    print_info_main(f"Loaded null embeds: {cfg.project.null_pe} {tuple(null_prompt_embeds.shape)}, pooled {tuple(null_pooled_prompt_embeds.shape)}")

    vae_config_shift_factor = vae.config.shift_factor
    vae_config_scaling_factor = vae.config.scaling_factor
    print_info_main(f"VAE shift factor: {vae_config_shift_factor}, scaling factor: {vae_config_scaling_factor}")
    if mode.has_conditioning:
        logger.info(f"Training with additional conditioning channels: {condition_channels} (total transformer input = base + cond)")
    else:
        logger.info("Training in unconditional mode (no extra conditioning channels).")

    # Scheduler and training steps (iterable dataset): must use max_train_steps only.
    if cfg.train.max_train_steps is None:
        raise ValueError("For iterable datasets, set cfg.train.max_train_steps; num_train_epochs is not used.")

    # get_scheduler does not handle num_cycles for cosine scheduler, so apply manually if needed
    if cfg.scheduler.lr_scheduler == "cosine":
        lr_scheduler = get_cosine_schedule_with_warmup(
            optimizer=optimizer,
            num_warmup_steps=cfg.scheduler.lr_warmup_steps,
            num_training_steps=cfg.train.max_train_steps,
            num_cycles=cfg.scheduler.lr_num_cycles,
        )
    else:
        lr_scheduler = get_scheduler(
            cfg.scheduler.lr_scheduler,
            optimizer=optimizer,
            num_warmup_steps=cfg.scheduler.lr_warmup_steps,
            num_training_steps=cfg.train.max_train_steps,
            num_cycles=cfg.scheduler.lr_num_cycles,
            power=cfg.scheduler.lr_power,
        )

    # Prepare everything with our `accelerator`.
    transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
        transformer, optimizer, train_dataloader, lr_scheduler
    )

    # Named-parameter helper for sharded EMA.
    # Cached: param names and references are stable throughout training
    # (optimizer.step modifies .data in-place, Parameter objects unchanged).
    _named_trainable_cache: list[tuple[str, torch.nn.Parameter]] | None = None
    def _named_trainable():
        """Return (name, param) for all trainable model params (cached)."""
        nonlocal _named_trainable_cache
        if _named_trainable_cache is None:
            _named_trainable_cache = [
                (n, p) for n, p in unwrap_compiled_module(transformer).named_parameters()
                if p.requires_grad
            ]
        return _named_trainable_cache

    if use_ema:
        import torch.distributed as _dist
        _ws = _dist.get_world_size() if _dist.is_initialized() else 1
        _rk = _dist.get_rank() if _dist.is_initialized() else 0
        ema_model = ShardedEMA(
            _named_trainable(),
            decay=float(getattr(cfg.train, 'ema_decay', 0.9999)),
            device=accelerator.device,
            world_size=_ws,
            rank=_rk,
        )
        logger.info(f"EMA enabled (sharded): {ema_model}")

    register_matdiff_checkpoint_hooks(
        accelerator,
        transformer,
        transformer_cls=type(unwrap_compiled_module(transformer)),
        pretrained_model_name_or_path=cfg.model.pretrained_model_name_or_path,
        transformer_lora_config=transformer_lora_config,
        is_full_finetune=is_full_finetune,
        keep_trainable_fp32=keep_trainable_fp32,
        upcast_before_saving=cfg.train.upcast_before_saving,
        use_ema=use_ema,
        ema_model=ema_model,
        logger=logger,
    )

    # Initialize the trackers
    if accelerator.is_main_process:
        tracker_name = "matdiff_release"
        _run_name = os.path.basename(os.path.normpath(str(cfg.project.output_dir)))
        init_kwargs = {"wandb": {"name": _run_name}}
        accelerator.init_trackers(
            tracker_name,
            config=OmegaConf.to_container(cfg, resolve=True),
            init_kwargs=init_kwargs,
        )

    # Train!
    total_batch_size = cfg.data.train.batch_size * accelerator.num_processes * cfg.train.gradient_accumulation_steps
    effective_total_batch_size = total_batch_size * cfg.model.num_views

    logger.info("***** Running training *****")
    logger.info(f"  Views per group (num_views V) = {cfg.model.num_views}")
    logger.info(f"  Instantaneous group batch size per device (bD) = {cfg.data.train.batch_size}")
    logger.info(f"  Instantaneous effective batch per device (bDV) = {cfg.data.train.batch_size * cfg.model.num_views}")
    logger.info(f"  Total train batch size (bD * parallel * accum) = {total_batch_size}")
    logger.info(f"  Total effective train batch size (bDV * parallel * accum) = {effective_total_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {cfg.train.gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {cfg.train.max_train_steps}")
    logger.info(f"  Number of trainable parameters = {sum(p.numel() for p in transformer_lora_parameters)}")
    conditioning_dropout = cfg.train.conditioning_dropout
    logger.info(
        "  Conditioning dropout probabilities: "
        f"text_only={conditioning_dropout.text_only * 100}%, "
        f"control_only={conditioning_dropout.control_only * 100}%, "
        f"text_and_control={conditioning_dropout.text_and_control * 100}%"
    )
    global_step = 0
    checkpoint_policy = CheckpointPolicy(
        cfg.project.output_dir,
        latest_alias="latest",
        keep=cfg.train.checkpoints_total_limit,
        search_in_checkpoint_dir=True,
        missing_latest_ok=True,
        prune_before_save=True,
        reserve_slots_before_save=1,
    )

    # Potentially load in the weights and states from a previous save
    if cfg.train.resume_from_checkpoint:
        resume_request = str(cfg.train.resume_from_checkpoint)
        resume_path = checkpoint_policy.resolve(resume_request)

        if resume_path is None:
            accelerator.print(
                f"Checkpoint '{cfg.train.resume_from_checkpoint}' does not exist. Starting a new training run."
            )
            cfg.train.resume_from_checkpoint = None
            initial_global_step = 0
        else:
            accelerator.print(f"Resuming from checkpoint {resume_path.name}")
            # Load to CPU first to avoid transient VRAM spike from optimizer
            # state double-allocation (torch.load on GPU + deepcopy inside
            # load_state_dict).  States are moved to GPU by load_state_dict.
            accelerator.load_state(str(resume_path), map_location="cpu")
            global_step = checkpoint_step(resume_path)

            initial_global_step = global_step

    else:
        initial_global_step = 0

    progress_bar = tqdm(
        range(0, cfg.train.max_train_steps),
        initial=initial_global_step,
        desc="Steps",
        disable=not accelerator.is_local_main_process,
    )

    scheduler_sigmas = noise_scheduler.sigmas.to(accelerator.device)
    scheduler_timesteps = noise_scheduler.timesteps.to(accelerator.device)

    # Step-0 validation before training.
    if accelerator.is_main_process and val_data is not None:
        logger.info("Running step-0 validation to verify pretrained knowledge...")
        _run_validation_step(0, val_data)
        logger.info("Step-0 validation complete.")

    transformer.train()
    _data_iter = iter(train_dataloader)
    memory_logging_steps = int(getattr(cfg.train, "memory_logging_steps", 50))
    image_logging_steps = int(getattr(cfg.train, "image_logging_steps", 1000))
    latent_verification_steps = int(getattr(cfg.train, "latent_verification_steps", 10))
    perf_tracker = StepPerformanceTracker(
        accelerator=accelerator,
        rate_multipliers={
            "samples_per_sec": total_batch_size,
            "view_samples_per_sec": effective_total_batch_size,
        },
        memory_log_every=memory_logging_steps,
    )
    perf_tracker.start(global_step)

    while global_step < cfg.train.max_train_steps:
        try:
            batch = next(_data_iter)
        except StopIteration:
            _data_iter = iter(train_dataloader)
            continue
        with accelerator.accumulate(transformer):
            training_batch = prepare_matdiff_training_batch(
                batch,
                device=accelerator.device,
                dtype=weight_dtype,
                use_dual_branch=use_dual_branch,
                text_only_dropout_prob=conditioning_dropout.text_only,
                control_only_dropout_prob=conditioning_dropout.control_only,
                text_and_control_dropout_prob=conditioning_dropout.text_and_control,
                null_prompt_embeds=null_prompt_embeds,
                null_pooled_prompt_embeds=null_pooled_prompt_embeds,
            )
            geometry = training_batch.geometry
            model_input = training_batch.model_input
            cond_latents = training_batch.cond_latents
            prompt_embeds = training_batch.prompt_embeds
            pooled_prompt_embeds = training_batch.pooled_prompt_embeds

            # Sample noise that we'll add to the latents
            noise = torch.randn_like(model_input)

            # ╭───────────────────────────── Group-level timestep sampling ───────────────────────────╮
            # │ We sample one timestep per base group (b) and broadcast it across                     │
            # │ its branches (D) and views (V).                                                       │
            # │ This ensures consistent noise levels for the same scene across different branches.      │
            # ╰───────────────────────────────────────────────────────────────────────────────────────╯
            u_group = sample_timestep_density(
                weighting_scheme=cfg.weighting.weighting_scheme,
                batch_size=geometry.groups,
                logit_mean=cfg.weighting.logit_mean,
                logit_std=cfg.weighting.logit_std,
                mode_scale=cfg.weighting.mode_scale,
                logit_uniform_mix=getattr(cfg.weighting, "logit_uniform_mix", 0.0),
            )
            num_t = noise_scheduler.config.num_train_timesteps
            indices_group = (u_group * num_t).long().clamp_(0, num_t - 1).to(device=model_input.device)
            indices = geometry.repeat_groups(indices_group)  # (b*D*V,)
            timesteps = scheduler_timesteps[indices]

            # Add noise according to flow matching.
            # zt = (1 - texp) * x + texp * z1
            sigmas = lookup_sigmas(scheduler_sigmas, indices, n_dim=model_input.ndim, dtype=model_input.dtype)
            noisy_latents = (1.0 - sigmas) * model_input + sigmas * noise

            # Append conditioning channels (slice to configured size) to noisy input
            transformer_input = append_condition_channels(noisy_latents, cond_latents, condition_channels)

            attach_corr_to_batch(
                batch,
                corr_dilate_iterations=corr_dilate_iterations,
                use_global_pos=mode.use_global_pos,
                use_dual_branch=use_dual_branch,
                use_weaving_attention=use_weaving_attention,
                gradual_nheads=gradual_nheads,
                device=accelerator.device,
            )

            joint_kwargs = {}
            if use_caa:
                joint_kwargs["attn_payload"] = payload_to_device(
                    batch['attn_payload'], accelerator.device, pos_dtype=weight_dtype)

            model_pred = transformer(
                hidden_states=transformer_input,
                timestep=timesteps,
                encoder_hidden_states=prompt_embeds,
                pooled_projections=pooled_prompt_embeds,
                joint_attention_kwargs=joint_kwargs if joint_kwargs else None,
                return_dict=False,
            )[0]

            # Follow: Section 5 of https://huggingface.co/papers/2206.00364.
            # Preconditioning of the model outputs.
            model_pred = apply_flow_preconditioning(
                model_pred,
                noisy_latents,
                sigmas,
                enabled=cfg.weighting.precondition_outputs,
            )

            # these weighting schemes use a uniform timestep sampling
            # and instead post-weight the loss
            weighting = compute_loss_weighting_for_sd3(
                weighting_scheme=cfg.weighting.weighting_scheme,
                sigmas=sigmas,
            )

            target = flow_matching_target(
                model_input,
                noise,
                precondition_outputs=cfg.weighting.precondition_outputs,
            )

            # Compute regular loss.
            loss = weighted_flow_mse(model_pred, target, weighting)

            accelerator.backward(loss)
            if accelerator.sync_gradients:
                accelerator.clip_grad_norm_(transformer_lora_parameters, cfg.train.max_grad_norm)

            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()

            # EMA update (local shard only — zero communication)
            if use_ema and accelerator.sync_gradients:
                ema_model.step(_named_trainable())

        # Training batch visualization (raw images only — no encode/decode roundtrip)
        should_log_images = image_logging_steps > 0 and (
            global_step in {0, 1} or global_step % image_logging_steps == 0
        )
        if should_log_images:
            try:
                pixel_values_vis = batch.get("pixel_values")  # B, V, C, H, W
                pixel_values_vis = rearrange(pixel_values_vis, 'B V ... -> (B V) ...')
                log_training_batch(
                    accelerator=accelerator,
                    pixel_values=pixel_values_vis,
                    step=global_step,
                    max_batch=6,
                )
            except Exception:
                logger.warning(
                    "Unable to log training batch at step %s. Skipping.",
                    global_step,
                    exc_info=True,
                )

        # Latent decode verification: log original vs decoded-from-precomputed-latent
        # for the first 10 steps to verify precomputed latent quality
        if latent_verification_steps > 0 and global_step < latent_verification_steps:
            try:
                pixel_values_flat = rearrange(batch["pixel_values"], 'B V ... -> (B V) ...')
                cond_values_flat = rearrange(batch["cond_values"], 'B V ... -> (B V) ...')
                log_latent_verification(
                    accelerator=accelerator,
                    vae=vae,
                    matvae=matvae,
                    pixel_values=pixel_values_flat.to(dtype=weight_dtype),
                    model_input=model_input,
                    step=global_step,
                    vae_scaling_factor=vae_config_scaling_factor,
                    vae_shift_factor=vae_config_shift_factor,
                    use_dual_branch=use_dual_branch,
                    cond_values=cond_values_flat.to(dtype=weight_dtype),
                    num_views=geometry.views,
                    max_batch=6,
                )
            except Exception:
                logger.warning(
                    "Unable to log latent verification at step %s. Skipping.",
                    global_step,
                    exc_info=True,
                )

        if accelerator.sync_gradients:
            progress_bar.update(1)
            global_step += 1
        
            accelerator.wait_for_everyone()
            if global_step % cfg.train.checkpointing_steps == 0:
                if accelerator.is_main_process or accelerator.distributed_type == DistributedType.DEEPSPEED:
                    if accelerator.is_main_process:
                        checkpoint_policy.prune(reserve_slots=checkpoint_policy.reserve_slots_before_save)

                    save_path = checkpoint_policy.path_for_step(global_step)
                    accelerator.save_state(str(save_path))
                    logger.info(f"Saved state to {save_path}")

                # EMA: gather all shards and save inference-compatible file
                # (collective — all ranks participate in gather, rank 0 writes)
                accelerator.wait_for_everyone()
                if use_ema:
                    ema_model.save_checkpoint(str(checkpoint_policy.path_for_step(global_step)))
                    
                    
            accelerator.wait_for_everyone()

            # Validation (EMA gather runs on all ranks; only main runs inference)
            _need_val = (
                validation_enabled
                and (global_step % cfg.train.validation_steps == 0)
            )
            if _need_val:
                _nt = _named_trainable()
                if use_ema:
                    ema_model.store(_nt)
                    ema_model.copy_to(_nt)  # collective gather
                if accelerator.is_main_process:
                    _run_validation_step(global_step, val_data)
                if use_ema:
                    ema_model.restore(_nt)
                accelerator.wait_for_everyone()

            logs = {
                "loss": loss.detach().item(),
                "lr": lr_scheduler.get_last_lr()[0],
            }
            perf_logs, reset_memory = perf_tracker.collect(global_step)
            logs.update(perf_logs)
            progress_bar.set_postfix(
                loss=logs["loss"],
                lr=logs["lr"],
                steps_s=logs["train/steps_per_sec"],
                gb=logs.get("train/peak_allocated_gb", 0.0),
            )
            accelerator.log(logs, step=global_step)
            perf_tracker.reset(global_step, reset_memory=reset_memory)
            accelerator.wait_for_everyone()

        if global_step >= cfg.train.max_train_steps:
            break

    # Final save: regular model and EMA.
    accelerator.wait_for_everyone()
    save_final_model = bool(getattr(cfg.train, "save_final_model", True))
    final_validation_enabled = matdiff_final_validation_enabled(cfg.train, validation_enabled=validation_enabled)

    # Gather EMA state (collective — all ranks must participate)
    ema_state = None
    if use_ema and (save_final_model or final_validation_enabled):
        ema_state = ema_model.state_dict()  # collective gather → CPU dict
        logger.info("Gathered EMA state for final save/validation.")

    if accelerator.is_main_process:
        model_to_save = unwrap_compiled_module(transformer)

        if save_final_model:
            save_dtype = torch.float32 if cfg.train.upcast_before_saving else weight_dtype
            save_matdiff_transformer_state(
                model_to_save,
                cfg.project.output_dir,
                is_full_finetune=is_full_finetune,
                save_dtype=save_dtype,
            )
            logger.info("Saved regular model.")

            if save_matdiff_ema_state(ema_state, cfg.project.output_dir, save_dtype=save_dtype):
                logger.info("Saved EMA weights (ema_weights.safetensors)")
        else:
            logger.info("Skipping final model save because train.save_final_model=false.")

        if final_validation_enabled and ema_state is not None:
            copy_state_to_named_parameters(model_to_save, ema_state)

        if final_validation_enabled and val_data is not None:
            _run_validation_step(global_step, val_data, is_final=True)

    accelerator.end_training()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Train MatDiff model",
        epilog="Additional config overrides: model.num_views=8 train.batch_size=4"
    )
    parser.add_argument(
        "--config",
        type=str,
        default="configs/matdiff/default.yaml",
        help="Path to config file (supports base-variant inheritance)"
    )
    
    args, remaining = parser.parse_known_args()
    cfg = load_config(args.config, cli_overrides=remaining)
    os.makedirs(cfg.project.output_dir, exist_ok=True)
    OmegaConf.save(cfg, os.path.join(cfg.project.output_dir, "config.yaml"))

    main(cfg)