lipforcing/trainer.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from typing import Dict, Any, List, Optional, Callable
import os
import time
import gc

import torch
from torch.utils.data import DataLoader

from lipforcing.methods import FastGenModel
from lipforcing.utils.basic_utils import set_random_seed, set_tmp_random_seed
from lipforcing.utils.checkpointer import Checkpointer, FSDPCheckpointer
import lipforcing.utils.logging_utils as logger
from lipforcing.configs.config import BaseConfig
from lipforcing.callbacks.callback import CallbackDict
from lipforcing.utils import instantiate, basic_utils
import lipforcing.utils.distributed.ddp as ddp
import lipforcing.utils.distributed.fsdp as fsdp
from lipforcing.utils.distributed import synchronize, is_rank0, world_size
import torch.distributed as dist
from lipforcing.utils import set_global_vars, set_temp_global_vars
from lipforcing.utils import global_vars
from lipforcing.utils.autoresume import AutoResumeInterface, create_auto_resume


class Trainer:
    def __init__(self, config: BaseConfig, auto_resume: Optional[AutoResumeInterface] = None):
        """
        Initialize the Trainer.

        Args:
            config (BaseConfig): LipForcing config
            auto_resume (Optional[AutoResumeInterface]): Custom auto-resume implementation.
                If None, defaults to NoOpAutoResume (auto-resume disabled).
                See lipforcing.utils.autoresume for the interface and examples.
        """

        self.config = config
        set_global_vars(self.config.trainer.global_vars)

        # Initialize auto-resume (defaults to NoOpAutoResume if not provided)
        self.auto_resume = create_auto_resume(auto_resume)
        logger.info(f"Auto-resume: {type(self.auto_resume).__name__}")

        # Set random seed
        set_random_seed(config.trainer.seed, by_rank=True)

        # Initialize the callback functions.
        logger.info("Initializing callbacks (including wandb)...")
        self.callbacks = CallbackDict(config=config, trainer=self)
        logger.success("Callbacks initialized successfully")

        # Synchronize after callback initialization to handle wandb timing differences
        synchronize()
        logger.info("Callback synchronization complete")

        # Initialize the checkpointer.
        logger.info("Initializing checkpointer...")
        if self.config.trainer.fsdp:
            self.checkpointer = FSDPCheckpointer(self.config.trainer.checkpointer)
        else:
            self.checkpointer = Checkpointer(self.config.trainer.checkpointer)
        logger.success("Checkpointer initialized successfully")

    def run(
        self,
        model: FastGenModel,
    ) -> None:
        """
        Run the training loop

        Args:
            model (FastGenModel): Distillation model.
        """
        logger.info("Starting training")

        iter_start = 0
        logger.info("Initializing callbacks and model ...")
        self.callbacks.on_model_init_start(model)
        if self.config.trainer.checkpointer.pretrained_ckpt_path:
            # This typically only affects the first job in the auto-resume chain
            self.load_pretrained_ckpt(model)

        if self.config.trainer.fsdp and (
            self.config.model.precision_amp is not None
            and self.config.model.precision_amp != self.config.model.precision
        ):
            logger.warning(
                f"Autocast to {self.config.model.precision_amp} is enabled and FSDP is enabled. "
                f"While this is possible, it is not recommended."
            )

        logger.info("Starting model.on_train_begin ...")
        synchronize()
        model.on_train_begin(is_fsdp=self.config.trainer.fsdp)
        synchronize()
        logger.info("model.on_train_begin completed")

        # wrap model into DDP or FSDP
        assert not (
            self.config.trainer.ddp and self.config.trainer.fsdp
        ), "Model cannot be wrapped into both DDP and FSDP"
        if self.config.trainer.ddp:
            logger.info("Wrapping model into ddp ..")
            model_ddp = ddp.model_to_ddp(model)
            logger.info("DDP wrapping completed")
        elif self.config.trainer.fsdp:
            logger.info("Wrapping model into fsdp ..")
            model_ddp = fsdp.model_to_fsdp(
                model,
                min_num_params=self.config.trainer.fsdp_min_num_params,
                apply_cpu_offload=self.config.trainer.fsdp_cpu_offload,
                sync_module_states=self.config.model.fsdp_meta_init,
                sharding_group_size=self.config.trainer.fsdp_sharding_group_size,
            )
            logger.info("FSDP wrapping completed")
        else:
            model_ddp = model
        self.callbacks.on_model_init_end(model_ddp)
        synchronize()

        self.callbacks.on_optimizer_init_start(model)
        model.init_optimizers()
        self.callbacks.on_optimizer_init_end(model)

        self.callbacks.on_load_checkpoint_start(model)

        # Check if we are resuming from an auto-resume checkpoint
        self.auto_resume.init()
        auto_resume_details = self.auto_resume.get_resume_details()
        logger.info(f"Auto-Resume Details: {auto_resume_details}")
        autoresume_ckpt = auto_resume_details["save_path"] if auto_resume_details else None

        if self.config.trainer.resume or autoresume_ckpt is not None:
            logger.info("Loading checkpoints for resuming ..")

            # load previous checkpoint
            iter_start = self.checkpointer.load(
                model.model_dict,
                optimizer_dict=model.optimizer_dict,
                scheduler_dict=model.scheduler_dict,
                grad_scaler=model.grad_scaler,
                callbacks=self.callbacks,
                path=autoresume_ckpt,
                device=model.device,
            )
        self.callbacks.on_load_checkpoint_end(model, iteration=iter_start)

        # re-seed based on the current iteration for resuming
        set_random_seed(self.config.trainer.seed, iteration=iter_start, by_rank=True)

        # resume samplers and initiate the dataloaders
        self.callbacks.on_dataloader_init_start(model, iteration=iter_start)
        nimg = (
            iter_start * self.config.dataloader_train.batch_size * self.config.trainer.grad_accum_rounds * world_size()
        )
        for loader in ["dataloader_train", "dataloader_val"]:
            dataloader_config = getattr(self.config, loader, None)
            if getattr(dataloader_config, "sampler_start_idx", 0) is None:
                logger.info(f"Setting sampler start index to {nimg} images for {loader}")
                dataloader_config.sampler_start_idx = nimg

        logger.info("Instantiating dataloader...")
        dataloader_train = instantiate(self.config.dataloader_train)
        dataloader_val = (
            instantiate(self.config.dataloader_val) if getattr(self.config, "dataloader_val", None) else None
        )
        # Record on-the-fly encoding config (if any dataloader requests it) so
        # preprocess_data can lazily build the encoders in the main process.
        self._setup_on_the_fly_encoding(dataloader_train, dataloader_val)
        augment_pipe = instantiate(self.config.trainer.augment_pipe)
        self.callbacks.on_dataloader_init_end(model, dataloader_train, dataloader_val, iteration=iter_start)

        self.callbacks.on_train_begin(model, iteration=iter_start)
        logger.info(f"iter_start: {iter_start}")

        if iter_start == 0 and dataloader_val is not None and not getattr(
            self.config.trainer, "skip_initial_validation", False
        ):
            # validation before first training step
            self.validate(model_ddp, model, dataloader_val, iteration=iter_start)

        dataloader_train_iter = iter(dataloader_train)
        for iter_cur in range(iter_start + 1, self.config.trainer.max_iter):
            self.callbacks.on_training_step_begin(model, iteration=iter_cur)
            for grad_accum_iter in range(self.config.trainer.grad_accum_rounds):
                data = next(dataloader_train_iter)
                data = self.preprocess_data(model, data, augment_pipe)
                logger.debug(
                    f"iteration: {iter_cur} | grad_accum_iter: {grad_accum_iter} | data: {basic_utils.to_str(data)}"
                )

                # single training step
                self.callbacks.on_training_accum_step_begin(model, data, iteration=iter_cur, accum_iter=grad_accum_iter)
                loss_map, outputs = self.train_step(model_ddp, model, data, iter_cur, grad_accum_iter)

            self.callbacks.on_training_step_end(
                model=model,
                data_batch=data,
                output_batch=outputs,
                loss_dict=loss_map,
                iteration=iter_cur,
            )

            # save checkpoint (before validation so progress is not lost if validation is slow)
            just_saved_checkpoint = False
            latest_checkpoint_path = None
            if iter_cur % self.config.trainer.save_ckpt_iter == 0:
                latest_checkpoint_path = self.save_checkpoint(model, iter_cur)
                just_saved_checkpoint = True

            # validation
            if iter_cur % self.config.trainer.validation_iter == 0 and dataloader_val is not None:
                self.validate(model_ddp, model, dataloader_val, iteration=iter_cur)

            if self.auto_resume_exit(
                model, iter_cur, skip_if_just_saved=just_saved_checkpoint, recent_checkpoint_path=latest_checkpoint_path
            ):
                # termination requested
                self.callbacks.on_train_end(model, iteration=iter_cur)
                self.callbacks.on_app_end(model, iteration=iter_cur)
                logger.info("Taking a 10 sec nap and exiting training.")
                time.sleep(10)
                return

        logger.info("Training complete.")
        # validation in the end
        if dataloader_val is not None:
            self.validate(model_ddp, model, dataloader_val, iteration=self.config.trainer.max_iter)
        self.save_checkpoint(model, self.config.trainer.max_iter)
        self.callbacks.on_train_end(model, iteration=self.config.trainer.max_iter)
        self.callbacks.on_app_end(model, iteration=self.config.trainer.max_iter)
        logger.info("Taking a 10 sec nap and exiting training.")
        time.sleep(10)

    def load_pretrained_ckpt(self, model: FastGenModel, device: torch.device | str = "cpu"):
        """
        Load pretrained model weights from a checkpoint.
        """
        key_map = self.config.trainer.checkpointer.pretrained_ckpt_key_map
        # use FSDP checkpointer to load the pretrained checkpoint
        # (which falls back to the basic checkpointer if the checkpoint ends with .pth)
        _checkpointer = FSDPCheckpointer(self.config.trainer.checkpointer)
        resume_iter = None
        for k_model, k_ckpt in key_map.items():
            if hasattr(model, k_model):
                model_dict = torch.nn.ModuleDict({k_ckpt: getattr(model, k_model)})
                resume_iter = _checkpointer.load(
                    model_dict,
                    path=self.config.trainer.checkpointer.pretrained_ckpt_path,
                    device=device,
                )
                logger.info(
                    f"Loaded {k_model} model from {k_ckpt} in {self.config.trainer.checkpointer.pretrained_ckpt_path} "
                    f"at iteration {resume_iter}"
                )
            else:
                logger.warning(
                    f"Model does not have submodule {k_model}. Skipping loading {k_ckpt} from "
                    f"{self.config.trainer.checkpointer.pretrained_ckpt_path}."
                )
        if resume_iter is not None:
            logger.info(f"Setting resume_iter for model to {resume_iter}.")
            model.resume_iter = resume_iter

    def save_checkpoint(self, model: FastGenModel, iteration: int, path: str | None = None) -> str:
        logger.info(f"Saving checkpoint iteration {iteration}")
        self.callbacks.on_save_checkpoint_start(model, iteration=iteration)
        # awaken the dataloader to avoid timeout
        path = self.checkpointer.save(
            model.model_dict,
            optimizer_dict=model.optimizer_dict,
            scheduler_dict=model.scheduler_dict,
            grad_scaler=model.grad_scaler,
            callbacks=self.callbacks,
            path=path,
            iteration=iteration,
        )
        self.callbacks.on_save_checkpoint_success(model, iteration=iteration, path=path)
        # Explicitly clear memory after checkpointing: we need this to
        # avoid OOM during wandb logging where the VAE is loaded and
        # used for decoding
        gc.collect()
        torch.cuda.empty_cache()
        self.callbacks.on_save_checkpoint_end(model, iteration=iteration)
        return path

    def train_step(
        self,
        model_ddp: FastGenModel | torch.nn.parallel.DistributedDataParallel,
        model: FastGenModel,
        data: Dict[str, Any],
        iteration: int,
        grad_accum_iter: int,
    ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
        """
        Single training step

        Args:
            model_ddp (FastGenModel | torch.nn.parallel.DistributedDataParallel): Distillation model with ddp wraaper.
            model (FastGenModel): Distillation model.
            data (Dict[str, Any]): Data dict for the current iteration.
            iteration: Current training iteration
            grad_accum_iter (int): Gradient accumulation iteration

        Returns:
            loss_map (dict[str, torch.Tensor]): Dictionary containing the loss values
            outputs (dict[str, torch.Tensor]): Dictionary containing the network output
        """
        grad_accum_rounds = self.config.trainer.grad_accum_rounds
        sync_grads = grad_accum_iter == grad_accum_rounds - 1

        if not self.config.trainer.fsdp:
            with ddp.ddp_sync_grad(model_ddp, sync_grads):
                # forward pass
                with model.autocast():
                    loss_map, outputs = model_ddp.single_train_step(data, iteration)
                # backward pass
                self.callbacks.on_backward_begin(
                    model, data, outputs, loss_map, iteration=iteration, accum_iter=grad_accum_iter
                )
                model.grad_scaler.scale(loss_map["total_loss"] / grad_accum_rounds).backward()
        else:
            with fsdp.fsdp_sync_grad(model, sync_grads):
                # forward pass
                with model.autocast():
                    loss_map, outputs = model_ddp.single_train_step(data, iteration)
                # backward pass
                self.callbacks.on_backward_begin(
                    model, data, outputs, loss_map, iteration=iteration, accum_iter=grad_accum_iter
                )
                model.grad_scaler.scale(loss_map["total_loss"] / grad_accum_rounds).backward()

        if grad_accum_iter == grad_accum_rounds - 1:
            # optimizer step, scheduler step, and more
            self.callbacks.on_optimizer_step_begin(model=model, iteration=iteration)
            model.optimizers_schedulers_step(iteration)
            # Zero after step to free memory on active optimizers
            model.optimizers_zero_grad(iteration)

        # detach loss_map and outputs
        return basic_utils.detach(loss_map), basic_utils.detach(outputs)

    @torch.no_grad()
    def validate(
        self,
        model_ddp: FastGenModel | torch.nn.parallel.DistributedDataParallel,
        model: FastGenModel,
        dataloader_val: DataLoader,
        iteration: int = 0,
    ) -> None:
        for idx, val_vars in enumerate(self.config.trainer.global_vars_val):
            with set_temp_global_vars(val_vars), set_tmp_random_seed(
                self.config.trainer.val_seed,
                by_rank=True,
                devices=[model.device] if model.device.type == "cuda" else [],
            ):
                self.callbacks.on_validation_begin(model, iteration=iteration, idx=idx)
                logger.info(f"Validation iteration {iteration}")
                for step, data in enumerate(dataloader_val):
                    if getattr(global_vars, "MAX_VAL_STEPS", None) is not None and step >= getattr(
                        global_vars, "MAX_VAL_STEPS"
                    ):
                        break

                    self.callbacks.on_validation_step_begin(model, data, step=step, iteration=iteration, idx=idx)
                    data = self.preprocess_data(model, data)
                    logger.debug(f"[val] step {step}: data preprocessed")
                    with model.autocast():
                        # Use validation_step if available (causal AR inference),
                        # otherwise fall back to single_train_step
                        if hasattr(model, "validation_step"):
                            loss_map, outputs = model.validation_step(data, iteration)
                        else:
                            loss_map, outputs = model_ddp.single_train_step(data, iteration)
                    self.callbacks.on_validation_step_end(
                        model, data, outputs, loss_map, step=step, iteration=iteration, idx=idx
                    )
                self.callbacks.on_validation_end(model, iteration=iteration, idx=idx)
                synchronize()

    def _setup_on_the_fly_encoding(self, *dataloaders) -> None:
        """Record on-the-fly encoder config from whichever dataloader requests it.

        The encoders themselves are built lazily on the first raw batch (in the
        main process), so precompute-only runs never touch them.
        """
        self._onfly_cfg = None
        self._onfly_encoder = None
        self._onfly_text_cache = {}
        for dl in dataloaders:
            if dl is not None and getattr(dl, "on_the_fly", False):
                cfg = dict(getattr(dl, "encoder_config", {}) or {})
                cfg["num_video_frames"] = getattr(dl, "num_video_frames", 81)
                cfg["cache_encoded"] = getattr(dl, "cache_encoded", True)
                cfg["cache_dir"] = getattr(dl, "cache_dir", None)
                self._onfly_cfg = cfg
                logger.info(
                    "On-the-fly preprocessing enabled; encoders built lazily in the "
                    f"main process (vae={cfg.get('vae_path')}, wav2vec={cfg.get('wav2vec_path')}, "
                    f"text_encoder={cfg.get('text_encoder_path')})."
                )
                break

    @torch.no_grad()
    def _maybe_encode_on_the_fly(self, model: FastGenModel, data: Dict[str, Any]) -> Dict[str, Any]:
        """Encode a raw on-the-fly batch into model-ready tensors (main process).

        No-op for fast-path (precomputed) batches. Lazily builds the frozen
        encoders on first use.
        """
        if not (isinstance(data, dict) and "_samples" in data):
            return data

        from lipforcing import preprocess as pp
        from lipforcing.datasets.omniavatar_dataloader import encode_on_the_fly_batch

        cfg = getattr(self, "_onfly_cfg", None) or {}
        if getattr(self, "_onfly_encoder", None) is None:
            vae_path = (cfg.get("vae_path")
                        or getattr(self.config.model, "vae_path", None))
            wav2vec_path = cfg.get("wav2vec_path")
            text_encoder_path = cfg.get("text_encoder_path")
            assert vae_path and wav2vec_path, (
                "On-the-fly encoding requires vae_path and wav2vec_path. Set them on the "
                "dataloader config (vae_path/wav2vec_path/text_encoder_path) or "
                "config.model.vae_path."
            )
            logger.info("Building on-the-fly encoders (frozen, eval) in the main process ...")
            self._onfly_encoder = pp.load_encoders(
                vae_path, wav2vec_path, model.device, dtype=model.precision,
                text_encoder_path=text_encoder_path,
                load_text=text_encoder_path is not None,
            )

        cache_dir = cfg.get("cache_dir")
        text_cache_dir = os.path.join(cache_dir, "_text_emb") if cache_dir else None
        return encode_on_the_fly_batch(
            self._onfly_encoder, data,
            num_video_frames=cfg.get("num_video_frames", 81),
            device=model.device, dtype=model.precision,
            cache_encoded=cfg.get("cache_encoded", True),
            text_emb_cache=self._onfly_text_cache,
            text_cache_dir=text_cache_dir,
        )

    @torch.no_grad()
    def preprocess_data(
        self, model: FastGenModel, data: Dict[str, Any], augment_pipe: Optional[Callable] = None
    ) -> Dict[str, Any]:
        """
        Preprocess the data before passing to the model.
        Args:
            model: FastGenModel
            data: Dict[str, Any]

        Returns:
            Dict[str, Any]: Preprocessed data

        """
        # On-the-fly: encode raw batches into model-ready tensors before anything
        # else (no-op for precomputed fast-path batches).
        data = self._maybe_encode_on_the_fly(model, data)
        ctx = dict(device=model.device, dtype=model.precision)
        data = basic_utils.to(data, **ctx)
        if augment_pipe is not None:
            data = augment_pipe(data)

        # we do not use torch.inference_mode here since resulting inference tensors
        # give runtime errors with gradient computations/checkpointing
        with torch.autocast(
            device_type=model.device.type, dtype=model.precision_amp_enc, enabled=model.precision_amp_enc is not None
        ):
            # Data/noise
            raw = "{}_raw".format
            for k in ["real", "noise"]:
                if k in data and raw(k) not in data:
                    data[raw(k)] = data[k]
                    # dataloader returns real of shape [B, C, T, H, W]
                    if hasattr(model.net, "vae") and data[k].shape[1] != model.input_shape[0]:
                        # Encode the data/noise to latent space
                        data[k] = model.net.vae.encode(data[k])

            # Text conditions
            for k in ["condition", "neg_condition"]:
                if k in data and raw(k) not in data:
                    data[raw(k)] = data[k]
                    if hasattr(model.net, "text_encoder") and isinstance(data[k], List):
                        # Encode the prompt to embedding
                        data[k] = model.net.text_encoder.encode(data[k])

            # Context for i2v/vid2vid
            if "real_raw" in data:
                if getattr(model.net, "is_i2v", False):  # extra vid context for i2v
                    # compute input for I2V models
                    real_raw_first_frame = data["real_raw"][:, :, 0:1]
                    bsz, channels = real_raw_first_frame.shape[0:2]
                    num_frames, height, width = data["real_raw"].shape[2:]
                    first_frame_cond = real_raw_first_frame

                    # Wan 2.1 I2V model concatenates first_frame_cond with noisy latents and mask.
                    # Wan 2.2 5B model replaces the first noisy latent frame with the first clean latent frame.
                    if model.net.concat_mask:
                        padding_shape = (bsz, channels, num_frames - 1, height, width)
                        first_frame_cond = torch.cat(
                            [real_raw_first_frame, real_raw_first_frame.new_zeros(*padding_shape)], dim=2
                        )

                    if hasattr(model.net, "vae"):
                        # Official Wan I2V implementation uses the VAE encoder with "argmax" mode to avoid stochasticity
                        data["first_frame_cond"] = model.net.vae.encode(first_frame_cond, mode="argmax")
                    else:
                        data["first_frame_cond"] = first_frame_cond

                if hasattr(model.net, "image_encoder"):
                    # Encode the first video frame with CLIP
                    data["encoder_hidden_states_image"] = model.net.image_encoder.encode(data["real_raw"][:, :, 0])

                if getattr(model.net, "is_vid2vid", False):  # extra vid context for vid2vid
                    assert hasattr(
                        model.net, "prepare_vid_conditioning"
                    ), "model.net must have prepare_vid_conditioning method"
                    if "depth_latent" in data:
                        data["vid_context"] = model.net.prepare_vid_conditioning(
                            data["real_raw"], condition_latents=data["depth_latent"]
                        )
                    else:
                        data["vid_context"] = model.net.prepare_vid_conditioning(data["real_raw"])

                # Cosmos video2world conditioning: use first a few frames as conditioning
                if getattr(model.net, "is_video2world", False):
                    num_cond_frames = getattr(model.net, "num_conditioning_frames", 1)
                    real_raw_first_frames = data["real_raw"][:, :, :num_cond_frames]
                    bsz, channels, _, height, width = data["real_raw"].shape

                    # Encode conditioning frames with VAE
                    if hasattr(model.net, "vae"):
                        data["conditioning_latents"] = model.net.vae.encode(real_raw_first_frames, mode="argmax")
                    else:
                        data["conditioning_latents"] = real_raw_first_frames

                    # Create condition mask: 1 for conditioning frames, 0 for generated frames
                    t_latent = data["real"].shape[2]
                    t_cond_latent = data["conditioning_latents"].shape[2]
                    condition_mask = torch.zeros(bsz, 1, t_latent, height // 8, width // 8, device=data["real"].device)
                    condition_mask[:, :, :t_cond_latent] = 1.0
                    data["condition_mask"] = condition_mask

        # Move encoded data to dtype and device
        data = basic_utils.to(data, **ctx)

        return data

    def auto_resume_exit(
        self, model: FastGenModel, iteration: int, skip_if_just_saved: bool = False, recent_checkpoint_path: str = None
    ) -> bool:
        """
        Check if the training should be terminated and auto-resume should be triggered.

        Args:
            model (FastGenModel): Distillation model.
            iteration (int): Current training iteration
            skip_if_just_saved (bool): Skip saving checkpoint if we just saved one
            recent_checkpoint_path (str): Path to the most recently saved checkpoint

        Returns:
            bool: True if the training should be terminated, False otherwise

        """
        # Check termination on rank 0 and broadcast to all ranks
        termination_requested = False

        # Ensure all ranks are ready before rank 0 checks termination
        synchronize()

        if is_rank0():
            termination_requested = self.auto_resume.termination_requested()

        # Broadcast the decision from rank 0 to all other ranks
        if world_size() > 1:
            termination_tensor = torch.tensor([1.0 if termination_requested else 0.0], device=model.device)
            dist.broadcast(termination_tensor, src=0)
            termination_requested = termination_tensor.item() > 0.5

        # Ensure all ranks have received the broadcast before proceeding
        synchronize()

        if not termination_requested:
            return False

        # Termination requested - save checkpoint and request resume
        ar_details = self.auto_resume.get_resume_details() or {}

        # Only save checkpoint if we haven't just saved one
        if not skip_if_just_saved:
            save_path = self.save_checkpoint(
                model, iteration, path=os.path.join(self.config.trainer.checkpointer.save_dir, "latest_ar.pth")
            )
            ar_details["save_path"] = save_path
        else:
            # Use the most recent checkpoint path
            logger.info("Skipping AutoResume checkpoint save as we just saved a regular checkpoint")
            if recent_checkpoint_path:
                save_path = recent_checkpoint_path
                logger.info(f"Using recently saved checkpoint: {save_path}")
            else:
                # Fallback: construct the path (this should rarely happen)
                logger.warning("No recent checkpoint path provided, constructing path")
                if isinstance(self.checkpointer, FSDPCheckpointer):
                    save_path = os.path.join(self.config.trainer.checkpointer.save_dir, f"{iteration:07d}")
                else:
                    save_path = os.path.join(self.config.trainer.checkpointer.save_dir, f"{iteration:07d}.pth")
            ar_details["save_path"] = save_path

        if is_rank0():
            self.auto_resume.request_resume(user_dict=ar_details)

        logger.info("Autoresume requested. Terminating training.")
        return True