# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch
import os
import time
import random
import numpy as np
import json
import threading
import shutil
import glob
from typing import Optional, Dict
from transformers import AutoConfig, GenerationConfig

from safetensors.torch import save_file
from huggingface_hub import create_repo, upload_folder, whoami
from huggingface_hub.utils import disable_progress_bars, enable_progress_bars
from cosmos_rl.utils.s3_utils import upload_folder_to_s3

from cosmos_rl.utils.logging import logger
from cosmos_rl.policy.trainer.optm import build_optimizers
from cosmos_rl.policy.model import ModelRegistry
from cosmos_rl.policy.config import Config as CosmosConfig
from cosmos_rl.utils.parallelism import ParallelDims
import cosmos_rl.utils.util as util
from cosmos_rl.utils.fp8.fp8_util import FP8ModelConverter
from cosmos_rl.policy.kernel.modeling_utils import init_flash_attn_meta
from cosmos_rl.utils.activation_offloading import get_act_offloading_ctx_manager
from cosmos_rl.policy.trainer.base import Trainer
from cosmos_rl.policy.trainer.base import extract_from_cuda_tensor, wrap_to_cuda_tensor
from cosmos_rl.utils.checkpoint import CheckpointMananger
from cosmos_rl.dispatcher.data.packer.base import BaseDataPacker
from cosmos_rl.policy.trainer.optm import build_lr_schedulers
from functools import partial


class LLMTrainer(Trainer):
    def __init__(
        self,
        config: CosmosConfig,
        parallel_dims: ParallelDims,
        train_stream: torch.cuda.Stream,
        data_packer: BaseDataPacker,
        val_data_packer: BaseDataPacker,
        **kwargs,
    ):
        super(LLMTrainer, self).__init__(
            config,
            parallel_dims,
            train_stream=train_stream,
            data_packer=data_packer,
            val_data_packer=val_data_packer,
            **kwargs,
        )

        if config.train.seed:
            torch.manual_seed(config.train.seed)
            torch.cuda.manual_seed(config.train.seed)
            torch.cuda.manual_seed_all(config.train.seed)
            random.seed(config.train.seed)
            np.random.seed(config.train.seed)

        if config.train.deterministic:
            torch.backends.cudnn.deterministic = True
            torch.backends.cudnn.benchmark = False
            torch.use_deterministic_algorithms(mode=True, warn_only=True)

        init_flash_attn_meta(
            config.train.deterministic, config.train.compile, config.train.fa_version
        )

        self.hf_config = util.retry(AutoConfig.from_pretrained)(
            config.policy.model_name_or_path,
            trust_remote_code=True,
        )
        model = ModelRegistry.build_model(config)

        # FP8 settings
        with torch.device("meta"):
            if config.train.fp8.enable_fp8:
                self.model_converter = FP8ModelConverter(config, parallel_dims)
                self.model_converter.convert_model(model)
            elif config.train.fp4.enable_fp4:
                from cosmos_rl.utils.fp4.fp4_util import FP4ModelConverter

                self.model_converter = FP4ModelConverter(config, parallel_dims)
                self.model_converter.convert_model(model)

        if config.train.fsdp_offload:
            model._apply(
                lambda t: torch.empty_like(t, device="cpu")
                if t.device.type == "meta"
                else t.to("cpu"),
                recurse=True,
            )

        try:
            # Apply parallelism to the model
            parallelize_fn, _ = model.parallelize_fn
            # `pp_scheduler` is used for both `sft` and `RLHF`
            # `pp_scheduler_val` is used only for `sft`, since `RLHF` does not require policy model via validation
            self.pp_scheduler, self.pp_scheduler_val = parallelize_fn(
                model, parallel_dims, config, pp_loss_fn=self.pp_loss_fn
            )
            if not config.train.fsdp_offload:
                model._apply(
                    lambda t: torch.empty_like(t, device=self.device)
                    if t.device.type == "meta"
                    else t.to("cuda"),
                    recurse=True,
                )
            model.post_to_empty_hook(config)
            if config.policy.lora is not None:
                from cosmos_rl.policy.lora.plugin import reinitialize_lora_params

                reinitialize_lora_params(model)
            # Enable gradient checkpointing for the model
            model.set_gradient_checkpointing_enabled(
                config.policy.model_gradient_checkpointing
            )

            torch.cuda.empty_cache()

            # If learning rate is list of `Union[float, List[float]]`,
            #  it means the learning rates are in order of pre-defined `model.separate_model_parts()`.
            # else, learning rate must be of type `Dict[str, float]`, where the key is the module path of the model part defined in `model.separate_model_parts()`, and the value is the learning rate for that model part.

            if isinstance(config.train.optm_lr, (float, list)):
                self.model_parts = model.separate_model_parts()
                self.model_modpath = [None for _ in range(len(self.model_parts))]
                # `named_modules()` recursively iterates over all modules in the model, and returns both the name and the module itself. We can use it to find the module path for each model part, which will be used in optimizer construction to get the parameters of each model part.
                # `named_children()` only iterates over the immediate children modules, which is not enough for us to find the module path for each model part since the model part can be nested in the model.
                for name, module in model.named_modules():
                    if module in self.model_parts:
                        idx = self.model_parts.index(module)
                        self.model_modpath[idx] = name
                        if all([modpath is not None for modpath in self.model_modpath]):
                            break
                for i in range(len(self.model_modpath)):
                    if self.model_modpath[i] is None:
                        self.model_modpath[i] = f"model_part_{i}"
            else:
                assert isinstance(config.train.optm_lr, dict), (
                    "Learning rate must be either a float, a list of floats, or a dict of {model_part: lr}."
                )
                import re
                from typing import Tuple, Union, Iterable, List

                _num_re = re.compile(r"(\d+)")

                def natural_key(s: str) -> Tuple[Union[int, str], ...]:
                    # makes 'layers.2' < 'layers.10'
                    parts = _num_re.split(s)
                    out = []
                    for p in parts:
                        out.append(int(p) if p.isdigit() else p)
                    return tuple(out)

                def sort_module_paths_deep_to_root(
                    paths: Iterable[str], sep: str = "."
                ) -> List[str]:
                    def key(p: str):
                        depth = 0 if not p else p.count(sep) + 1
                        # primary: deeper first; secondary: natural lexicographic for stable sibling ordering
                        return (-depth, natural_key(p))

                    return sorted(paths, key=key)

                model_parts = []
                module_paths = []
                learning_rates = []

                # check for global learning rate if exists
                global_lr = config.train.optm_lr.get("global", None)
                if global_lr is not None:
                    del config.train.optm_lr["global"]
                    assert isinstance(global_lr, float) and global_lr > 0, (
                        f"Global learning rate must be a positive float if specified, but got {global_lr}."
                    )

                # Check existence of the module paths specified in the learning rate dict
                module_paths = sort_module_paths_deep_to_root(
                    config.train.optm_lr.keys()
                )
                all_named_modules = dict(model.named_modules())
                for module_path in module_paths:
                    if not self.parallel_dims.pp_enabled:
                        assert module_path in all_named_modules, (
                            f"Module path {module_path} specified in learning rate dict not found in the model. Available module paths are: {list(all_named_modules.keys())}"
                        )
                    elif module_path not in all_named_modules:
                        logger.warning(
                            f"Module path {module_path} specified in learning rate dict not found in the model. Available module paths are: {list(all_named_modules.keys())}"
                        )
                    else:
                        pass
                    model_parts.append(all_named_modules[module_path])
                    learning_rates.append(config.train.optm_lr[module_path])

                if global_lr is not None:
                    model_parts.append(model)
                    # empty string is used to represent the global learning rate for the whole model, which does not correspond to any module path.
                    module_paths.append("[ALL_OTHER]")
                    learning_rates.append(global_lr)

                self.model_parts = model_parts
                self.model_modpath = module_paths
                # reset the learning rate in config to be the list of learning rates for each model part, which will be used in optimizer construction.
                # because modules have to be in specific order during optimizer building,
                # either a list or an ordered dict is needed to ensure the order of the model parts and the learning rates.
                self.config.train.optm_lr = learning_rates

            self.model = model
            # util.add_nan_checks(model)
        except Exception as e:
            import traceback

            traceback.print_exc()
            raise e

        self.ckpt_manager = CheckpointMananger(
            self.config,
            self.parallel_dims,
            self.global_rank,
            hook_fns=kwargs.get("hook_fns", {}),
        )
        # FIXME: (lms) use_streams=True could cause NaN in backward. Fix this later.
        self.act_offloading_ctx_manager = get_act_offloading_ctx_manager(
            self.model, config.train.activation_offload, use_streams=False
        )

        self.build_optimizers()

        self.seq_len_multiple = parallel_dims.cp * parallel_dims.tp
        self.lr_schedulers = None
        self.upload_thread = None
        if self.config.train.fp8.enable_fp8 or self.config.train.fp4.enable_fp4:
            # Constraint of FP8/FP4 kernel(torch._scaled_mm): it requires K in MNK is mutiple of 16. In backward of Linear, to
            # calculate the gradient of weight, we have to round the seq_len_multiple to mutiple of 16.
            # See: https://github.com/pytorch/pytorch/blob/851a6fa82df251fbc368f0284d941ce78f68e7b1/aten/src/ATen/native/cuda/Blas.cpp#L1252
            self.seq_len_multiple = (self.seq_len_multiple + 16 - 1) // 16 * 16
            logger.info(
                "FP8/FP4 Training enabled, round seq_len_multiple to mutiple of 16."
            )
        logger.info(
            f"Trainer initialized at local rank {self.local_rank}, with seq_len_multiple: {self.seq_len_multiple}"
        )

    def build_optimizers(self):
        # TODO(cjx): add `CompiledAutograd` support
        self.optimizers = build_optimizers(
            self.model_parts, self.config, model_modpath=self.model_modpath
        )
        if self.config.train.fp8.enable_fp8 or self.config.train.fp4.enable_fp4:
            self.optimizers.register_step_post_hook(
                lambda *args, **kwargs: self.model_converter.post_optimizer_hook(
                    self.model_parts
                )
            )

    def sync_all_states(
        self,
        is_send: bool,
        send_hook: callable,
        recv_hook: callable,
        has_reference_model: bool = False,
    ) -> int:
        """
        Sync all states of the model and optimizer.
        Args:
            is_send (bool): Whether to send or receive the states.
            send_hook (callable): The hook function to send the states.
            recv_hook (callable): The hook function to receive the states.
            reference_model (bool): Whether to sync the reference model state dict.
        Returns:
            len_params (int): The number of parameters synced.
        """
        len_params = 0
        # It's a HFModel, we need to sync the named buffers
        state_dict = self.model.state_dict()
        model_state_dict = [state_dict]

        if has_reference_model:
            if len(self.reference_state_dict) == 0:
                assert not is_send, (
                    "Reference model state dict should be populated before sending"
                )
                for key, value in model_state_dict[0].items():
                    self.reference_state_dict[key] = torch.empty_like(
                        value, device="cpu"
                    )
            model_state_dict.append(self.reference_state_dict)
        model_state_dict[0].update(dict(self.model.named_buffers()))

        # 1. Sync all model states
        for state_to_sync in model_state_dict:
            for dest_name in sorted(state_to_sync.keys()):
                obj = state_to_sync[dest_name]
                assert isinstance(obj, torch.Tensor)
                local_view = wrap_to_cuda_tensor(
                    self.device, dest_name, obj, in_place=obj.is_cuda
                )
                if is_send:
                    send_hook(local_view)
                else:
                    recv_hook(local_view)
                    if isinstance(obj, torch.distributed.tensor.DTensor):
                        to_write = obj.to_local()
                    else:
                        to_write = obj

                    # Copy again for offloaded tensor since it is not inplace received
                    if not to_write.is_cuda:
                        to_write.copy_(local_view)
                len_params += 1

        # 2. Sync optimizer states
        optimizer_state = self.optimizers.state_dict()
        for dest_name in sorted(optimizer_state.keys()):
            obj = optimizer_state[dest_name]
            local_view = wrap_to_cuda_tensor(self.device, dest_name, obj)
            if local_view.data_ptr() is None:
                # skip the optimizer state if the data pointer is None
                continue
            if is_send:
                # nccl send
                send_hook(local_view)
            else:
                # nccl recv
                recv_hook(local_view)
                optimizer_state[dest_name] = extract_from_cuda_tensor(
                    self.device,
                    dest_name,
                    obj,
                    local_view,
                )
            len_params += 1

        if not is_send:
            self.optimizers.load_state_dict(optimizer_state)

        # 3. Sync lr_scheduler states
        if self.lr_schedulers is not None:
            lr_sheduler_state = self.lr_schedulers.state_dict()
            for dest_name in sorted(lr_sheduler_state.keys()):
                obj = lr_sheduler_state[dest_name]
                local_view = wrap_to_cuda_tensor(self.device, dest_name, obj)
                if is_send:
                    # nccl send
                    send_hook(local_view)
                else:
                    # nccl recv
                    recv_hook(local_view)
                    lr_sheduler_state[dest_name] = extract_from_cuda_tensor(
                        self.device,
                        dest_name,
                        obj,
                        local_view,
                    )
                len_params += 1
            if not is_send:
                self.lr_schedulers.load_state_dict(lr_sheduler_state)

        # 4. Sync rng_state
        rng_state = self.ckpt_manager.get_rng_state()
        for dest_name in sorted(rng_state.keys()):
            obj = rng_state[dest_name]
            local_view = wrap_to_cuda_tensor(self.device, dest_name, obj)
            if is_send:
                # nccl send
                send_hook(local_view)
            else:
                # nccl recv
                recv_hook(local_view)
                rng_state[dest_name] = extract_from_cuda_tensor(
                    self.device, dest_name, obj, local_view
                )
            len_params += 1
        if not is_send:
            self.ckpt_manager.set_rng_state(rng_state)
        return len_params

    def export_safetensors(
        self,
        output_dir: str,
        rel_path: str,
        trainable_only: bool = False,
        is_final=False,
        dtype: Optional[torch.dtype] = None,
    ):
        """Export safetensors to the local directory/huggingface/s3.
        Args:
            output_dir (str): The directory to save the safetensors.
            rel_path (str): The relative path to the output directory.
            trainable_only (bool): Whether to export only the trainable parameters.
            is_final (bool): Whether this is the final export.
            dtype (torch.dtype): The dtype of the parameters.
        """
        if self.config.policy.lora is not None and not trainable_only:
            trainable_only = True
            logger.info(
                "Exporting safetensors with param `trainable_only` is overridden to `True` for LoRA."
            )

        save_hf_config = self.config.policy.lora is None
        save_lora_config = self.config.policy.lora is not None
        save_generation_config = True
        export_weight_index_json = self.config.policy.lora is None

        path = os.path.join(output_dir, rel_path)
        if self.parallel_dims.dp_replicate_coord[0] > 0:
            return

        if self.global_rank == 0:
            logger.info(
                f"Prepare to exporting safetensors to {path} at rank {self.global_rank}"
            )

        if not self.parallel_dims.dp_replicate_enabled:
            torch.distributed.barrier()

        def get_tensor_size(tensor):
            """Get the size of the tensor in bytes."""
            return tensor.element_size() * tensor.numel()

        max_file_size_gb = 4 if not save_lora_config else float("inf")
        max_size_bytes = max_file_size_gb * 1024**3  # 4 GB in bytes
        current_chunk = {}
        total_chunk_size = 0
        current_chunk_size = 0
        file_idx = 0
        manifest = {}  # Record the weight->file name mapping
        cpu_chunks_to_save = []

        def create_file_name(save_lora_config, pp_rank, pp_size, file_idx):
            if save_lora_config:
                name = "adapter_model.safetensors"
            else:
                if pp_size == 1:
                    name = f"{file_idx:05d}.safetensors"
                else:
                    name = f"model-{pp_rank}-of-{pp_size}-{file_idx:05d}.safetensors"
            return os.path.join(path, name)

        def save_chunked_tensors(
            chunk: Dict[str, torch.Tensor], chunk_size: int, file_path: str
        ):
            """
            Save a dictionary of tensors into a safetensors file.
            Only the rank 0 of dp_shard, cp, tp will save the file.
            """
            nonlocal total_chunk_size
            if (
                self.parallel_dims.dp_shard_coord[0] == 0
                and self.parallel_dims.cp_coord[0] == 0
                and self.parallel_dims.tp_coord[0] == 0
            ):
                os.makedirs(os.path.dirname(file_path), exist_ok=True)
                chunk_cpu = {}
                for name, param in chunk.items():
                    manifest[name] = os.path.basename(file_path)
                    chunk_cpu[name] = param.cpu()
                total_chunk_size += chunk_size
                cpu_chunks_to_save.append(
                    (chunk_cpu, file_path)
                )  # Save for CPU saving in upload handler
                logger.info(
                    f"Chunk {file_idx} to be saved at {os.path.basename(file_path)}"
                )

        for name, param in self.model.named_parameters():
            # First map the key from local to hf naming convention
            name = self.model.weight_mapper.policy_map_local_key_to_hf_key(name)
            if trainable_only and not param.requires_grad:
                continue
            is_dtensor = isinstance(param, torch.distributed.tensor.DTensor)
            param = param.full_tensor() if is_dtensor else param
            param = param.detach().data

            pp_rank, pp_size = self.parallel_dims.pp_coord

            for (
                _name,
                _param,
            ) in self.model.weight_mapper.policy_map_local_key_for_export_tensor(
                name, param
            ):
                if _param is None:
                    logger.debug(
                        f"[Policy] Skipping None parameter for {name} in safetensors export."
                    )
                    continue
                elif save_lora_config and not _name.startswith("base_model"):
                    # LoRA model needs to add a prefix to the weight name to be consistent with the HF naming convention
                    _name = f"base_model.model.{_name}"

                _param = _param.to(dtype=dtype) if dtype is not None else _param
                tensor_size = get_tensor_size(_param)
                # If adding the current tensor exceeds the size limit, save the current chunk
                if current_chunk_size + tensor_size > max_size_bytes:
                    # Save the current chunk as a safetensor file
                    file_name = create_file_name(
                        save_lora_config, pp_rank, pp_size, file_idx
                    )
                    save_chunked_tensors(current_chunk, current_chunk_size, file_name)

                    # Reset for the next chunk
                    current_chunk = {_name: _param}
                    current_chunk_size = tensor_size
                    file_idx += 1
                else:
                    # Add tensor to the current chunk
                    current_chunk[_name] = _param
                    current_chunk_size += tensor_size

        # Save any remaining tensors in the last chunk
        if current_chunk:
            file_name = create_file_name(save_lora_config, pp_rank, pp_size, file_idx)
            save_chunked_tensors(current_chunk, current_chunk_size, file_name)

        # Allgather the manifest from all pipeline stages
        if self.parallel_dims.pp_enabled:
            pp_group = self.parallel_dims.mesh["pp"].get_group()
            pp_size = self.parallel_dims.pp
            output = [None for _ in range(pp_size)]
            tensor_sizes = [0 for _ in range(pp_size)]
            torch.distributed.all_gather_object(output, manifest, group=pp_group)
            torch.distributed.all_gather_object(
                tensor_sizes, total_chunk_size, group=pp_group
            )
            merged_manifest = {}
            for m in output:
                merged_manifest.update(m)
            total_tensor_size = sum(tensor_sizes)
        else:
            merged_manifest = manifest
            total_tensor_size = total_chunk_size

        if not self.parallel_dims.dp_replicate_enabled:
            torch.distributed.barrier()

        def save_and_upload_handler(
            chunks_to_save,
            weight_index_json,
            hf_config,
            lora_config,
            data_packer,
            generation_config,
            config,
            is_final,
            path,
            rel_path,
            max_retries=3,
        ):
            """Handle the save to local and the upload of the model to huggingface and s3."""
            # upload the final model to huggingface
            for chunk_cpu, file_path in chunks_to_save:
                save_file(chunk_cpu, file_path)
            if weight_index_json is not None:
                with open(os.path.join(path, "model.safetensors.index.json"), "w") as f:
                    json.dump(
                        weight_index_json,
                        f,
                        indent=4,
                    )
            # save hf_config
            if hf_config is not None:
                hf_config.save_pretrained(path)
            if lora_config is not None:
                # Save the LoRA config
                with open(os.path.join(path, "adapter_config.json"), "w") as f:
                    json.dump(lora_config, f, indent=4)
            if data_packer is not None:
                data_packer.save_state(path)
            # save the generation config to get the generation aligned with HF.
            if generation_config is not None:
                generation_config.save_pretrained(path)

            # Copy missing .py and .json files from source model directory
            # This is needed for models loaded with trust_remote_code=True
            # Only copy files that don't exist in destination to avoid overwriting
            try:
                from cosmos_rl.utils.util import resolve_model_path

                source_model_path = resolve_model_path(
                    config.policy.model_name_or_path,
                    revision=config.policy.model_revision,
                )

                # Find all .py and .json files in source
                source_files = []
                source_files.extend(glob.glob(os.path.join(source_model_path, "*.py")))
                source_files.extend(
                    glob.glob(os.path.join(source_model_path, "*.json"))
                )

                copied_files = []
                for src_file in source_files:
                    filename = os.path.basename(src_file)
                    dst_file = os.path.join(path, filename)

                    # Only copy if file doesn't exist in destination
                    if not os.path.exists(dst_file):
                        shutil.copy2(src_file, dst_file)
                        copied_files.append(filename)

                if copied_files:
                    logger.info(
                        f"Copied missing files from source model: {copied_files}"
                    )
            except Exception as e:
                logger.warning(f"Failed to copy missing files from source model: {e}")

            if config.train.ckpt.upload_hf and is_final:
                username = whoami()["name"]
                repo_id = (
                    username
                    + "/"
                    + config.train.ckpt.hf_repo_name
                    + "-"
                    + config.train.timestamp
                )
                logger.info(f"Uploading the final model to huggingface: {repo_id}...")
                retry = 0
                success = False
                while retry < max_retries:
                    try:
                        create_repo(repo_id, exist_ok=True)
                        # hide redundant logs of huggingface
                        disable_progress_bars()
                        upload_folder(
                            folder_path=path,
                            path_in_repo=".",
                            repo_id=repo_id,
                            commit_message="Upload model",
                        )
                        enable_progress_bars()
                        logger.info(f"Model uploaded to huggingface: {repo_id}")
                        success = True
                        break
                    except Exception as e:
                        logger.error(f"Failed to upload model to huggingface: {e}")
                        retry += 1
                if not success:
                    logger.error(
                        "All retry attempts to upload model to huggingface failed."
                    )
                    raise RuntimeError(
                        f"Failed to upload model to huggingface after {max_retries} attempts."
                    )
            # upload the model to s3
            if config.train.ckpt.upload_s3:
                if is_final:
                    # syncronizely upload the final model to s3
                    upload_folder_to_s3(
                        path,
                        config.train.ckpt.s3_bucket,
                        os.path.join(config.train.ckpt.s3_prefix, rel_path),
                    )
                elif config.train.ckpt.upload_s3 == "all":
                    # asynchronously upload the model to s3
                    upload_folder_to_s3(
                        path,
                        config.train.ckpt.s3_bucket,
                        os.path.join(config.train.ckpt.s3_prefix, rel_path),
                    )
            logger.info(f"\n\nExported safetensors to {path}\n\n")

        if self.global_rank == 0:
            if export_weight_index_json:
                weight_index_json = {
                    "metadata": {
                        "total_size": total_tensor_size,
                    },
                    "weight_map": merged_manifest,
                }
            else:
                weight_index_json = None

            # save hf_config
            if save_hf_config:
                hf_config = self.hf_config
            else:
                hf_config = None

            if save_lora_config:
                # Save the LoRA config
                lora_config = self.config.policy.lora.model_dump(mode="json")
                lora_config["base_model_name_or_path"] = (
                    self.config.policy.model_name_or_path
                )
                lora_config["peft_type"] = "LORA"
            else:
                lora_config = None

            # save the generation config to get the generation aligned with HF.
            generation_config = None
            try:
                if save_generation_config:
                    generation_config = util.retry(GenerationConfig.from_pretrained)(
                        self.config.policy.model_name_or_path
                    )
            except Exception:
                logger.warning("[Policy] No generation config found, do not save it.")

            # If the upload thread is already running, wait for it to finish
            if self.upload_thread is not None:
                self.upload_thread.join()
            self.upload_thread = threading.Thread(
                target=save_and_upload_handler,
                args=(
                    cpu_chunks_to_save,
                    weight_index_json,
                    hf_config,
                    lora_config,
                    self.data_packer,
                    generation_config,
                    self.config,
                    is_final,
                    path,
                    rel_path,
                ),
                name="save_and_upload_safetensors",
                daemon=True,
            )
            self.upload_thread.start()

    def model_load_from_hf(self):
        start_time = time.time()
        self.model.load_hf_weights(
            self.config.policy.model_safetensor_path
            or self.config.policy.model_name_or_path,
            self.parallel_dims,
            self.device,
            revision=self.config.policy.model_revision,
        )
        end_time = time.time()
        logger.info(
            f"Time taken to load model from HF: {end_time - start_time:.2f} seconds"
        )

    def model_resume_from_checkpoint(self):
        ckpt_extra_vars, self.lr_schedulers = self.ckpt_manager.load_checkpoint(
            model=self.model,
            optimizer=self.optimizers,
            scheduler=partial(build_lr_schedulers, self.optimizers, self.config),
            model_name_or_path=self.config.policy.model_name_or_path,
            revision=self.config.policy.model_revision,
        )
        return ckpt_extra_vars

    def step_validation(self):
        # By default, add an empty step_validation for LLM trainer.
        # For GRPO, validation is handled in rollout side, so this method is not needed.
        # For SFT trainer, we will override this method in sft_trainer.py.
        pass
