VLA/wall_x/trainer/qwen_vl_act_trainer.py

import os
import gc
import time
import torch
import random
import numpy as np
import torch.nn as nn
import torch.distributed as dist

from tqdm import tqdm
from functools import wraps
from datetime import datetime
from torch.optim import AdamW
from torch.distributed.tensor import distribute_tensor
from accelerate import Accelerator
from safetensors.torch import load_file
from transformers.optimization import get_cosine_with_min_lr_schedule_with_warmup
from transformers import AutoProcessor
from wall_x.utils.timers import Timers
from wall_x.model.qwen2_5_based import Qwen2_5_VLMoEForAction, Qwen2_5_VLConfig
from wall_x.data.config import ACTION_DATASET_NAMES, MULTIMODAL_DATASET_NAMES
from wall_x.data.load_lerobot_dataset import (
    PreprocessedDataset,
    get_data_configs,
    load_lerobot_data,
)


def timer(func):
    """
    Decorator to measure function execution time.

    Args:
        func: Function to be timed

    Returns:
        Wrapped function with timing functionality
    """

    @wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        print(
            f"\033[92m[current time: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}] Function {func.__name__} took {end_time - start_time:.2f} seconds to execute\033[0m"
        )
        return result

    return wrapper


def print_rank_last(message):
    """
    Print message only on the last rank in distributed training.

    Args:
        message (str): Message to print
    """
    if torch.distributed.is_initialized():
        if torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1):
            print(message, flush=True)
    else:
        print(message, flush=True)


def seed_all(seed):
    """
    Set random seeds for reproducible training.

    Args:
        seed (int): Random seed value
    """
    np.random.seed(seed)
    random.seed(seed)
    torch.manual_seed(seed)


class QwenVlAct_Trainer:
    """
    Vision-Language-Action trainer for Qwen-VL models with robotic action prediction.

    This trainer handles multi-modal learning combining vision, language, and action data
    for robotic control applications. It supports distributed training, mixed precision,
    gradient accumulation, and various optimization strategies including MoE (Mixture of Experts).

    Features:
    - Multi-modal data processing (vision + language + actions)
    - Distributed training with Accelerate
    - Gradient accumulation and clipping
    - Learning rate scheduling with warmup
    - Checkpoint saving and resuming
    - Comprehensive logging and monitoring
    """

    @timer
    def __init__(
        self,
        config,
        logger,
        accelerator: Accelerator = None,
        seed=42,
        data_config_path=None,
    ):
        """
        Initialize the Vision-Language-Action trainer.

        Args:
            config (dict): Training configuration dictionary containing:
                - processor_path (str): Path to data preprocessing processor
                - qwen_vl_act_config_path (str): Path to model configuration file
                - learning_rate (float): Base learning rate for training
                - num_epoch (int): Number of training epochs
                - pretrained_wallx_path (str): Path to pretrained model
                - And other training hyperparameters
            logger: Logger instance for tracking metrics
            accelerator (Accelerator, optional): Hugging Face Accelerate instance for distributed training
            seed (int, optional): Random seed for reproducibility. Defaults to 42.
            data_config_path (str, optional): Path to data configuration file

        Raises:
            ValueError: If required configuration keys are missing
        """
        # Validate required configuration keys
        required_keys = ["learning_rate", "num_epoch"]
        for key in required_keys:
            if key not in config:
                raise ValueError(f"Missing required configuration key: {key}")

        self.config = config
        self.logger = logger
        self.accelerator = accelerator
        self.seed = seed

        # Initialize random seeds for reproducibility
        seed_all(self.seed)

        # Training state variables
        self.start_epoch = 0
        self.global_step = 0
        self.num_epoch = self.config["num_epoch"]
        self.initial_step = 0

        # Data and model configuration
        self.dataload_config = get_data_configs(self.config["data"])
        self.data_config_path = data_config_path
        self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False)

        # Load model and initialize training components
        self.load_model()
        self.action_dim = sum(self.config["dof_config"].values())

        # Distributed training setup
        self.rank = self.accelerator.process_index
        self.world_size = self.accelerator.num_processes
        print(
            f"rank {self.accelerator.process_index} after load model memory usage: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB",
            flush=True,
        )

        # Load training data
        self.load_qact_data()
        print(
            f"rank {self.accelerator.process_index} after load qact data usage: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB",
            flush=True,
        )

        # Resume from checkpoint if specified
        if "resume" in self.config:
            self.resume_from_checkpoint()

        # Initialize special token IDs
        self.propri_token_id = self.processor.tokenizer.convert_tokens_to_ids(
            "<|propri|>"
        )
        self.action_token_id = self.processor.tokenizer.convert_tokens_to_ids(
            "<|action|>"
        )

        # Initialize evaluation metrics
        self.base_l1_loss = None
        self.base_l1_loss_detail = {}

        # Performance monitoring
        self.timers = Timers(log_level=0, log_option="minmax")

        # Adjust global step if resuming from checkpoint
        if self.initial_step != 0:
            self.global_step = self.initial_step // self.config.get(
                "gradient_accumulation_steps", 1
            )

    def print_rank0(self, msg, flush=True):
        """
        Print message only on rank 0 (main process).

        Args:
            msg: Message to print
            flush (bool): Whether to flush output buffer
        """
        if self.accelerator.is_main_process:
            print(msg, flush=flush)

    def fit(self):
        """
        Main training loop executing multiple epochs with validation.

        Handles the complete training process including:
        - Training loop execution
        - Validation after each epoch
        - Process synchronization
        - Memory cleanup
        """
        self.accelerator.wait_for_everyone()
        # Optional validation before training starts
        if self.config.get("resume", None) is not None and self.config["resume"].get(
            "validate_first", False
        ):
            self.val_loop()
            self.accelerator.wait_for_everyone()

        # Main training loop
        for epoch in range(self.start_epoch, self.num_epoch):
            self.train_loop(epoch)
            self.accelerator.wait_for_everyone()

            if (epoch + 1) % self.config.get("epoch_save_interval", 1) == 0:
                self.save_checkpoint(epoch)

            # Validation after each epoch
            self.val_loop()
            self.accelerator.wait_for_everyone()

            # Memory cleanup
            gc.collect()

    def train_loop(self, epoch):
        """
        Execute training for a single epoch.

        Args:
            epoch (int): Current epoch number

        Handles:
        - Data loading and batching
        - Forward/backward passes
        - Gradient accumulation and clipping
        - Learning rate scheduling
        - Loss logging and monitoring
        - Performance profiling (optional)
        """
        # Initialize training dataloader for current epoch
        if isinstance(self.dataset, PreprocessedDataset):
            if getattr(self, "train_dataloader", None) is not None:
                self.dataset._train()
                self.train_sampler.set_epoch(epoch)
            else:
                self.train_dataloader, self.train_sampler = (
                    self.dataset.get_train_dataloader()
                )
                self.train_sampler.set_epoch(epoch)
        else:
            self.train_dataloader = self.dataset.get_train_dataloader()

        self.model.train()
        grad_accum_steps = self.config.get("gradient_accumulation_steps", 1)
        total = len(self.train_dataloader)
        t0 = time.time()
        enable_profiling = self.config["profile"]

        # Optional PyTorch profiler for performance analysis
        if enable_profiling:
            profiler = torch.profiler.profile(
                activities=[
                    torch.profiler.ProfilerActivity.CPU,
                    torch.profiler.ProfilerActivity.CUDA,
                ],
                schedule=torch.profiler.schedule(
                    wait=self.config["profile_wait_iters"],
                    warmup=self.config["profile_warmup_iters"],
                    active=self.config["profile_active_iters"],
                ),
                on_trace_ready=torch.profiler.tensorboard_trace_handler(
                    self.config["profile_save_path"], worker_name="worker0"
                ),
                record_shapes=True,
                profile_memory=True,
                with_stack=True,
            )
            profiler.__enter__()

        try:

            # Setup timers for First iteration
            self.timers("interval-time", log_level=0).start(barrier=False)
            self.timers("data-load", log_level=0).start(barrier=False)

            for i, batch in enumerate(self.train_dataloader, self.initial_step):
                # Move batch to device
                if isinstance(self.dataset, PreprocessedDataset):
                    batch = {
                        k: (
                            v.to(self.accelerator.device, non_blocking=True)
                            if isinstance(v, torch.Tensor)
                            else v
                        )
                        for k, v in batch.items()
                    }

                self.timers("data-load").stop()

                with self.accelerator.accumulate(self.model):
                    # Forward pass
                    self.timers("forward-compute", log_level=0).start(barrier=False)
                    outputs = self.model(**batch, mode="train")
                    self.timers("forward-compute").stop()

                    loss = outputs.loss
                    # Check for NaN loss
                    if torch.isnan(loss):
                        print(
                            f"Warning: NaN loss detected in epoch: {epoch}, step: {i}",
                            flush=True,
                        )
                        continue

                    # Backward pass
                    self.timers("backward-compute", log_level=0).start(barrier=False)
                    self.accelerator.backward(loss)
                    self.timers("backward-compute").stop()

                    # Gradient clipping
                    total_norm = self.accelerator.clip_grad_norm_(
                        self.model.parameters(), self.config.get("max_grad_norm", 1.0)
                    )

                    # Optimizer step
                    self.timers("optimizer", log_level=0).start(barrier=False)
                    self.optimizer.step()
                    self.optimizer.zero_grad()
                    self.timers("optimizer").stop()

                    # Update global step and learning rate after gradient accumulation
                    if (i + 1) % grad_accum_steps == 0:
                        self.lr_scheduler.step()
                        self.global_step += 1
                        lr = self.lr_scheduler.get_last_lr()[0]

                        # Gather loss across all processes for logging
                        train_loss = (
                            self.accelerator.gather(loss.detach()).mean().item()
                        )
                        _log_dict = {
                            "lr": lr,
                            "train_loss": train_loss,
                        }

                        # Log component losses
                        if (
                            "cross_entropy_loss" in outputs
                            and outputs.cross_entropy_loss is not None
                        ):
                            _log_dict["cross_entropy_loss"] = (
                                self.accelerator.gather(
                                    outputs.cross_entropy_loss.detach()
                                )
                                .mean()
                                .item()
                            )

                        if "flow_loss" in outputs and outputs.flow_loss is not None:
                            _log_dict["flow_loss"] = (
                                self.accelerator.gather(outputs.flow_loss.detach())
                                .mean()
                                .item()
                            )

                        # Log per-dataset channel losses
                        if (
                            "channel_loss_dict" in outputs
                            and outputs.channel_loss_dict is not None
                        ):
                            for dataset_name_i in (
                                ACTION_DATASET_NAMES + MULTIMODAL_DATASET_NAMES
                            ):
                                count_sum = (
                                    self.accelerator.gather(
                                        outputs.channel_loss_count_dict[dataset_name_i]
                                    )
                                    .sum()
                                    .item()
                                )
                                if count_sum > 0:
                                    channel_loss = (
                                        self.accelerator.gather(
                                            outputs.channel_loss_dict[
                                                dataset_name_i
                                            ].detach()
                                        )
                                        .sum()
                                        .item()
                                        / count_sum
                                    )
                                    _log_dict[f"channel_loss_{dataset_name_i}"] = (
                                        channel_loss
                                    )

                            # Log action accuracy for fast tokenizer
                            if (
                                "action_accuracy" in outputs.channel_loss_dict
                                and self.use_fast_tokenizer
                            ):
                                _log_dict["action_accuracy"] = (
                                    self.accelerator.gather(
                                        outputs.channel_loss_dict[
                                            "action_accuracy"
                                        ].detach()
                                    )
                                    .mean()
                                    .item()
                                )

                        # Log metrics
                        if self.logger is not None:
                            self.logger.log(_log_dict, step=self.global_step)

                        # Log gradient norm
                        if self.logger is not None and self.accelerator.sync_gradients:
                            self.logger.log(
                                {"total_norm": total_norm}, step=self.global_step
                            )

                self.timers("interval-time").stop()

                # Setup timers for next iteration
                if i < len(self.train_dataloader) - 1:
                    self.timers("interval-time", log_level=0).start(barrier=False)
                    self.timers("data-load", log_level=0).start(barrier=False)

                # Periodic logging
                t1 = time.time()
                if i % 1 == 0:
                    lr = self.lr_scheduler.get_last_lr()[0]
                    self.training_log(
                        epoch, self.num_epoch, i, total, loss, lr, t1 - t0
                    )
                    t0 = time.time()

                if enable_profiling:
                    profiler.step()

        finally:
            if enable_profiling:
                profiler.__exit__(None, None, None)

    @torch.no_grad()
    def val_loop(self):
        """
        Execute validation loop with gradient computation disabled.

        Evaluates model performance on validation set and logs validation loss.
        """
        # Initialize validation dataloader
        if getattr(self, "val_dataloader", None) is not None:
            self.dataset._eval()
            self.val_sampler.set_epoch(0)
        else:
            self.val_dataloader, self.val_sampler = self.dataset.get_val_dataloader()
            self.val_sampler.set_epoch(0)

        self.model.eval()
        self.val_loss = 0

        # Validation loop
        for i, batch in enumerate(
            tqdm(
                self.val_dataloader,
                desc="Validating",
                total=len(self.val_dataloader),
                disable=not self.accelerator.is_main_process,
            )
        ):
            if isinstance(self.dataset, PreprocessedDataset):
                batch = {
                    k: (
                        v.to(self.accelerator.device, non_blocking=True)
                        if isinstance(v, torch.Tensor)
                        else v
                    )
                    for k, v in batch.items()
                }

            with torch.no_grad():
                outputs = self.model(**batch, mode="train")
                loss = outputs.loss
                self.val_loss += self.accelerator.gather(loss.detach()).mean().item()

        # Calculate average validation loss
        self.val_loss /= len(self.val_dataloader)

        # Log validation metrics
        if self.logger is not None:
            self.logger.log({"val_loss": self.val_loss}, step=self.global_step)

        self.model.train()

    @timer
    def load_model(self):
        """
        Load and configure the Vision-Language-Action model.

        Handles:
        - Model loading from pretrained weights
        - Processor initialization
        - Optimizer configuration (with support for different learning rates for different components)
        - Learning rate scheduler setup
        - Model preparation for distributed training
        """
        # Load pretrained model
        model_type = self.config.get("model_type", "qwen2_5")
        assert model_type in ["wall-oss", "qwen2_5"]
        if model_type == "wall-oss":
            model = Qwen2_5_VLMoEForAction.from_pretrained(
                self.config["pretrained_wallx_path"],
                **{"use_fast_tokenizer": self.use_fast_tokenizer},
            )
            self.processor = model.processor
            model = model.to(torch.bfloat16)
        elif model_type == "qwen2_5":

            config = Qwen2_5_VLConfig.from_pretrained(
                self.config["qwen_vl_act_config_path"]
            )
            flow_loss_weight = self.config.get("flow_loss_weight", 1.0)
            self.processor = AutoProcessor.from_pretrained(
                self.config["pretrained_wallx_path"], use_fast=True
            )
            if self.config.get("use_fast_tokenizer", False):
                action_tokenizer_path = self.config["action_tokenizer_path"]
                action_tokenizer = AutoProcessor.from_pretrained(
                    action_tokenizer_path, trust_remote_code=True
                )
                # process for use fast
                new_tokens = ["<|propri|>", "<|action|>"]
                new_tokens += [
                    f"<|action_token_{i}|>" for i in range(action_tokenizer.vocab_size)
                ]
                self.processor.tokenizer.add_tokens(new_tokens)
                begin_idx_token = "<|action_token_0|>"
                token_id = self.processor.tokenizer.convert_tokens_to_ids(
                    begin_idx_token
                )
                self.processor.tokenizer.init_kwargs["action_token_start_index"] = (
                    token_id
                )
                self.processor.tokenizer.init_kwargs["action_token_vocab_size"] = (
                    action_tokenizer.vocab_size
                )
                self.processor.action_processor = action_tokenizer
            model = Qwen2_5_VLMoEForAction(
                config,
                self.use_fast_tokenizer,
                self.processor,
                flow_loss_weight=flow_loss_weight,
            )

            model = model.to(torch.bfloat16)
            model = self.load_qwen_pretrain_weight(
                model, self.config["pretrained_wallx_path"]
            )
            model.resize_token_embeddings(len(self.processor.tokenizer))
            model = model.to(torch.bfloat16)
        else:
            raise NotImplementedError(f"Invalid model type: {model_type}")

        # Configure optimizer based on training strategy
        if "freeze_vlm" in self.config and self.config["freeze_vlm"]:
            print("Freezing VLM parameters, training only MoE experts", flush=True)
            moe_params = []
            for name, param in model.named_parameters():
                if "moe.experts.1." not in name:
                    param.requires_grad = False
                else:
                    moe_params.append(param)
            param_groups = [{"params": moe_params, "lr": self.config["learning_rate"]}]
            self.optimizer = AdamW(param_groups, weight_decay=0.1)

        elif "action_expert_learning_rate" in self.config:
            # Separate learning rates for VLM and action expert parameters
            moe_params = []
            vlm_params = []
            for name, param in model.named_parameters():
                if "moe.experts.1." in name:
                    moe_params.append(param)
                else:
                    vlm_params.append(param)

            # Configure parameter groups
            if self.config.get("train_action_expert_only", False):
                self.print_rank0("Training action expert only", flush=True)
                param_groups = [
                    {
                        "params": moe_params,
                        "lr": self.config["action_expert_learning_rate"],
                    }
                ]
            else:
                param_groups = [
                    {"params": vlm_params, "lr": self.config["learning_rate"]},
                    {
                        "params": moe_params,
                        "lr": self.config["action_expert_learning_rate"],
                    },
                ]

            self.optimizer = AdamW(param_groups, weight_decay=0.1)
            self.print_rank0(
                f"Setting MoE learning rate to {self.config['action_expert_learning_rate']}, "
                f"VLM learning rate to {self.config['learning_rate']}",
                flush=True,
            )
        else:
            # Standard optimizer configuration
            self.optimizer = AdamW(
                model.parameters(),
                lr=self.config["learning_rate"],
                weight_decay=0.1,
            )

        # Configure learning rate scheduler
        warmup_steps = self.config.get("num_warmup_steps", 0)
        num_training_steps = self.config.get("num_training_steps", 1000000000)
        min_lr = self.config.get("min_lr", 0.1 * self.config["learning_rate"])
        self.lr_scheduler = get_cosine_with_min_lr_schedule_with_warmup(
            optimizer=self.optimizer,
            num_warmup_steps=warmup_steps,
            num_training_steps=num_training_steps,
            min_lr=min_lr,
        )

        self.model = model

        # Enable gradient computation for embeddings
        if hasattr(model, "enable_input_require_grads"):
            self.model.enable_input_require_grads()
        else:

            def make_inputs_require_grad(module, input, output):
                output.requires_grad_(True)

            self.model.get_input_embeddings().register_forward_hook(
                make_inputs_require_grad
            )

        # Prepare model, optimizer, and scheduler for distributed training
        self.model, self.optimizer, self.lr_scheduler = self.accelerator.prepare(
            self.model, self.optimizer, self.lr_scheduler
        )

    @timer
    def load_qact_data(self):
        """
        Load and configure training data for Vision-Language-Action learning.

        Supports LeRobot dataset format and handles distributed data loading
        across multiple processes.
        """
        print(f"Loading Vision-Language-Action data from {__file__}")
        self.accelerator.wait_for_everyone()

        # Load LeRobot dataset
        self.dataset, self.train_num = load_lerobot_data(
            self.config,
            self.dataload_config.get("lerobot_config", {}),
            rank=self.rank,
            world_size=self.world_size,
        )

    @timer
    def load_qwen_pretrain_weight(self, model, pretrain_weight_path):
        """
        Load pretrained Qwen weights with MoE adaptation.

        Args:
            model: Model instance to load weights into
            pretrain_weight_path (str): Path to pretrained weight files

        Returns:
            Model with loaded pretrained weights

        Handles weight key renaming for MoE architecture compatibility.
        """
        # Load all safetensors files
        weight_files = sorted(
            [f for f in os.listdir(pretrain_weight_path) if f.endswith(".safetensors")]
        )
        merged_weights = {}

        # Merge weights from all files
        for weight_file in weight_files:
            file_path = os.path.join(pretrain_weight_path, weight_file)
            weights = load_file(file_path)
            merged_weights.update(weights)

        # Rename weights for MoE compatibility
        renamed_weights = {}
        for key, value in merged_weights.items():
            if (
                key.startswith("model.layers")
                and "mlp." in key
                and model.config.mlp_moe
            ):
                # Rename MLP weights for MoE structure
                layer_num = key.split(".layers.")[1].split(".mlp")[0]
                new_key = key.replace(
                    f"layers.{layer_num}.mlp.", f"layers.{layer_num}.moe.experts.0."
                )
                renamed_weights[new_key] = value
            elif (
                key.startswith("model.layers")
                and "self_attn." in key
                and model.config.attention_moe
            ):
                # Rename attention weights for MoE structure
                layer_num = key.split(".layers.")[1].split(".self_attn")[0]
                proj_types = ["q_proj", "k_proj", "v_proj", "o_proj"]
                for proj in proj_types:
                    if proj in key:
                        new_key = key.replace(
                            f"layers.{layer_num}.self_attn.{proj}",
                            f"layers.{layer_num}.self_attn.{proj}_experts.0",
                        )
                        renamed_weights[new_key] = value
                        break
            else:
                renamed_weights[key] = value

        # Load weights into model
        err = model.load_state_dict(renamed_weights, strict=False)
        self.print_rank0(f"Weight loading report: {err}", flush=True)
        if self.accelerator.is_main_process:
            self.print_rank0(f"Loaded pretrained weights from: {pretrain_weight_path}")

        return model

    def training_log(
        self,
        current_epoch,
        total_epoch,
        current_train_iter,
        total_train_iter,
        loss,
        lr,
        time_per_step,
    ):
        """
        Log training progress and performance metrics.

        Args:
            current_epoch (int): Current epoch number
            total_epoch (int): Total number of epochs
            current_train_iter (int): Current training iteration
            total_train_iter (int): Total iterations in epoch
            loss (torch.Tensor): Current loss value
            lr (float): Current learning rate
            time_per_step (float): Time taken for current step
        """
        timers_to_log = [
            "interval-time",
            "data-load",
            "forward-compute",
            "backward-compute",
            "optimizer",
        ]

        log_string = f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]"
        log_string += " epoch {:3d}/{:3d} |".format(current_epoch, total_epoch)
        log_string += " iter {:6d}/{:6d} |".format(current_train_iter, total_train_iter)
        log_string += " loss {:.6f} |".format(loss)
        log_string += " lr {:.6f} |".format(lr)
        log_string += " time_per_step_avg {:.6f}s |".format(time_per_step)

        print_rank_last(log_string)
        self.timers.log(timers_to_log, normalizer=1)

    def save_checkpoint(self, epoch, step=0):
        """
        Save training checkpoint.

        Args:
            epoch (int): Current epoch number
            step (int, optional): Current step number. Defaults to 0.

        Saves model state, optimizer state, and training progress information.
        """
        save_path = self.config["save_path"]
        if step == 0:
            ckpt_path = f"{save_path}/{epoch}"
        else:
            ckpt_path = f"{save_path}/{epoch}_{step}"

        # If FSDP SHARDED_STATE_DICT is used, please refer to the wall-x/workspace/README.md
        # merge checkpoint section to merge the weights into a single safetensors if needed.
        self.accelerator.save_state(ckpt_path)

        self.processor.save_pretrained(os.path.join(ckpt_path, "processor"))

        # Save current iteration steps for dataset resuming
        if step != 0:
            _rank = self.accelerator.process_index
            if isinstance(self.dataset, PreprocessedDataset):
                torch.save(
                    {"epoch": epoch, "step": step},
                    os.path.join(
                        ckpt_path, f"epoch_{epoch}_step_{step}_rank_{_rank}.pth"
                    ),
                )

    def resume_from_checkpoint(self):
        """
        Resume training from a saved checkpoint.

        Handles both full checkpoint loading and model-only loading based on configuration.
        """
        checkpoint_path = self.config["resume"]["ckpt"]

        if self.config.get("resume", {}).get("load_ckpt_only", False):
            if self.config.get("FSDP2", False):
                self._load_fsdp_state_dict_with_distribute_tensor()

            else:
                # Load only model weights
                ckpt_path = self.config["resume"]["ckpt"] + "/model.safetensors"
                state_dict = load_file(ckpt_path, device="cpu")

                # Add module prefix if needed for distributed training
                new_state_dict = {}
                for key in state_dict:
                    if not key.startswith("module."):
                        new_key = "module." + key
                    new_state_dict[new_key] = state_dict[key]

                self.model.load_state_dict(new_state_dict, strict=False)
        else:
            # Load full checkpoint including optimizer and scheduler states
            self.accelerator.load_state(checkpoint_path)

        self.print_rank0(f"Resumed from checkpoint: {checkpoint_path}")

    def _load_fsdp_state_dict_with_distribute_tensor(self):

        rank = dist.get_rank() if dist.is_initialized() else 0

        full_sd = load_file(
            self.config["resume"]["ckpt"] + "/model.safetensors", device="cpu"
        )
        meta_sharded_sd = self.model.state_dict()
        sharded_sd = {}

        def find_matching_key(target_key, available_keys):
            if target_key in available_keys:
                return target_key
            prefixed_key = f"_orig_mod.{target_key}"
            if prefixed_key in available_keys:
                return prefixed_key
            if target_key.startswith("_orig_mod."):
                unprefixed_key = target_key[len("_orig_mod.") :]
                if unprefixed_key in available_keys:
                    return unprefixed_key
            return None

        for param_name, full_tensor in full_sd.items():
            matching_key = find_matching_key(param_name, meta_sharded_sd.keys())
            if matching_key is None:
                if rank == 0:
                    print(
                        f"[Rank {rank}] Warning: Parameter not found:",
                        param_name,
                        flush=True,
                    )
                continue
            sharded_meta_param = meta_sharded_sd[matching_key]
            sharded_tensor = distribute_tensor(
                full_tensor,
                sharded_meta_param.device_mesh,
                sharded_meta_param.placements,
            )
            sharded_sd[matching_key] = nn.Parameter(sharded_tensor)

        self.model.load_state_dict(sharded_sd, assign=True, strict=False)

    def log_l1_details(self, all_label, all_pred, all_task, all_dof_mask):
        """
        Log detailed L1 loss metrics by degrees of freedom.

        Args:
            all_label (torch.Tensor): Ground truth action labels
            all_pred (torch.Tensor): Predicted actions
            all_task (list): Task identifiers
            all_dof_mask (torch.Tensor): Degrees of freedom mask

        Computes and logs L1 loss for each DOF component separately for detailed analysis.
        """
        all_task = all_task[: len(all_label)]

        # Apply DOF mask
        all_label = all_label * all_dof_mask
        all_pred = all_pred * all_dof_mask

        # Compute baseline L1 loss (predict mean action)
        if self.base_l1_loss is None:
            mean_action = all_label.mean(dim=0)
            self.base_l1_loss = nn.functional.l1_loss(all_label, mean_action)

        self.logger.log(
            {"base_l1_loss": self.base_l1_loss.item()}, step=self.global_step
        )

        # Log L1 loss for each DOF component
        start_idx = 0
        dof_config = self.config.get("dof_config", {})
        for dof in dof_config:
            end_idx = start_idx + dof_config[dof]
            dof_label = all_label[:, :, start_idx:end_idx]
            dof_pred = all_pred[:, :, start_idx:end_idx]
            dof_l1 = nn.functional.l1_loss(dof_pred, dof_label)

            self.print_rank0(f"DOF {dof}, L1 loss: {dof_l1.item()}", flush=True)
            self.logger.log(
                {f"detail/l1_loss_{dof}": dof_l1.item()}, step=self.global_step
            )

            start_idx = end_idx