cosmos-framework/cosmos_framework/callbacks/hf_export.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
"""HFExportCallback: export VLM DCP checkpoints to HuggingFace safetensors format.

Design notes
------------
- Hooks into ``on_save_checkpoint`` (called by DistributedCheckpointer.save() before I/O).
- All ranks participate in the weight-gather phase (DTensor.full_tensor() all-gathers).
- Rank 0 accumulates CPU tensors, writes shards, and uploads — other ranks exit early.
- File I/O and upload run in a background thread on rank 0 to avoid blocking training.
- Worker exceptions are stored in ``_worker_exception`` and re-raised on the next
  checkpoint or at train end, so failures are never silently swallowed.
- Controlled entirely via ``config.checkpoint.hf_export`` (HFExportConfig).

Phase 2+ note
-------------
Weight parameters are iterated via ``model.model.model.named_parameters()`` where
``model.model`` is the ``HFModel`` wrapper and ``model.model.model`` is the underlying
HuggingFace transformer.  Parameter names are already HF-native — no weight_mapper
remapping is required.
"""

import json
import os
import shutil
import threading
from typing import Any

import torch

from cosmos_framework.utils import log
from cosmos_framework.utils.callback import Callback
from cosmos_framework.utils.distributed import is_rank0

try:
    from safetensors.torch import save_file as _safetensors_save_file
except ImportError:
    _safetensors_save_file = None

try:
    from transformers import AutoTokenizer, GenerationConfig
except ImportError:
    AutoTokenizer = None
    GenerationConfig = None

# Map string dtype names (as stored in ParallelismConfig.precision) to torch dtypes.
_DTYPE_MAP: dict[str, torch.dtype] = {
    "float32": torch.float32,
    "float16": torch.float16,
    "bfloat16": torch.bfloat16,
    "float64": torch.float64,
}


def _upload_folder_to_s3(local_folder: str, bucket: str, s3_prefix: str, credential_path: str) -> None:
    """Upload every file under *local_folder* to ``s3://{bucket}/{s3_prefix}/...``.

    Uses the i4 ``easy_io`` S3 backend (Boto3Backend), which reads credentials from
    *credential_path*.  Files are uploaded as streaming transfers via boto3's
    ``upload_file()`` — the full shard is never loaded into memory.
    """
    from cosmos_framework.utils.easy_io import easy_io

    backend = easy_io.get_file_backend(
        backend_args={
            "backend": "s3",
            "s3_credential_path": credential_path,
            "path_mapping": None,
        }
    )
    for root, _, files in os.walk(local_folder):
        for fname in sorted(files):
            local_path = os.path.join(root, fname)
            rel = os.path.relpath(local_path, local_folder)
            s3_path = f"s3://{bucket}/{s3_prefix}/{rel}"
            # Pass the local path string so Boto3Backend uses upload_file() —
            # a streaming transfer that avoids reading the whole shard into memory.
            backend.put(local_path, s3_path)
            log.info(f"[HFExportCallback] Uploaded {local_path} → {s3_path}")


class HFExportCallback(Callback):
    """Export VLM weights to HuggingFace-compatible safetensors after each DCP checkpoint.

    Enabled / configured via ``config.checkpoint.hf_export`` (HFExportConfig).  Disabled
    by default — add this callback and set ``hf_export.enabled = True`` to activate.

    Exports written to::

        {job.path_local}/hf_exports/iter_{iteration:09d}/
            00000.safetensors
            ...
            model.safetensors.index.json
            config.json
            tokenizer.json   (if tokenizer can be loaded from model_name_or_path)

    Optionally uploads to:
    - S3 (``hf_export.upload_to_object_store``)
    - HuggingFace Hub (``hf_export.hf_repo_id``)

    Args:
        dtype: Export weight dtype (e.g. ``"bfloat16"``).  Use
            ``"${model.config.precision}"`` in the Hydra callback config to
            inherit from the training precision.
    """

    # HuggingFace convention: max 4 GB per shard file.
    _MAX_SHARD_BYTES: int = 4 * 1024**3

    def __init__(self, dtype: str = "bfloat16") -> None:
        self._export_dtype: torch.dtype | None = _DTYPE_MAP.get(dtype)
        self._current_iteration: int = 0
        self._export_thread: threading.Thread | None = None
        # Stores any exception raised inside the background worker so it can be
        # re-raised on the main thread at the next checkpoint or train end.
        self._worker_exception: BaseException | None = None

    # ------------------------------------------------------------------
    # Callback hooks
    # ------------------------------------------------------------------

    def on_save_checkpoint_start(self, model: Any, iteration: int = 0) -> None:
        self._current_iteration = iteration

    def on_save_checkpoint(self, model: Any, state_dict: dict[str, Any]) -> None:
        hf_cfg = self.config.checkpoint.hf_export
        if not hf_cfg.enabled:
            return

        iteration = self._current_iteration
        if iteration % hf_cfg.export_every_n != 0:
            return

        # Deferred import to avoid circular dependency at module load time.
        from cosmos_framework.model.vfm.vlm_model import VLMModel

        if not isinstance(model, VLMModel):
            # The legacy vlm/train.py path passes model_parts: list[nn.Module] (raw HF
            # models without the VLMModel attribute structure).  HF export requires the
            # VLMModel wrapper, which is only available via the unified cosmos_framework/scripts/train.py path.
            if isinstance(model, list):
                log.warning(
                    "[HFExportCallback] Received model_parts (list) instead of VLMModel. "
                    "HF export requires the unified training path (cosmos_framework/scripts/train.py). Skipping."
                )
            else:
                log.warning(
                    "[HFExportCallback] model is not VLMModel (got %s); skipping HF export.",
                    type(model).__name__,
                )
            return

        if _safetensors_save_file is None:
            raise ImportError("safetensors is required for HFExportCallback. Install it with: pip install safetensors")

        output_dir = os.path.join(self.config.job.path_local, "hf_exports", f"iter_{iteration:09d}")

        # ----------------------------------------------------------------
        # Phase 1 (all ranks): gather sharded parameters into CPU chunks.
        # full_tensor() is a collective operation — all ranks must participate.
        # ----------------------------------------------------------------
        cpu_chunks, manifest, total_size = self._gather_weights(model)

        # ----------------------------------------------------------------
        # Phase 2 (rank 0, background thread): file I/O + optional upload.
        # ----------------------------------------------------------------
        if not is_rank0():
            return

        # Block on any still-running export from the previous checkpoint and
        # propagate any worker exception before starting a new export.
        if self._export_thread is not None and self._export_thread.is_alive():
            log.warning(
                "[HFExportCallback] Previous export thread still running; waiting before starting export for iter %d.",
                iteration,
            )
            self._export_thread.join()

        if self._worker_exception is not None:
            exc = self._worker_exception
            self._worker_exception = None
            raise RuntimeError(f"[HFExportCallback] Previous export failed with: {exc}") from exc

        self._export_thread = threading.Thread(
            target=self._save_and_upload,
            args=(cpu_chunks, manifest, total_size, model.hf_config, model.model_name_or_path, output_dir, iteration),
            daemon=True,
        )
        self._export_thread.start()

    def on_train_end(self, model: Any, iteration: int = 0) -> None:
        """Wait for the final export thread so the process does not exit prematurely."""
        if self._export_thread is not None and self._export_thread.is_alive():
            log.info("[HFExportCallback] Waiting for export thread to finish...")
            self._export_thread.join()
            log.info("[HFExportCallback] Export thread done.")

        if self._worker_exception is not None:
            exc = self._worker_exception
            self._worker_exception = None
            raise RuntimeError(f"[HFExportCallback] Export thread failed with: {exc}") from exc

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _gather_weights(self, model: Any) -> tuple[list[dict[str, torch.Tensor]], dict[str, str], int]:
        """Iterate model parameters, all-gather DTensor shards, and build CPU chunks.

        Must be called on **all ranks**. Only rank 0 populates the returned
        ``cpu_chunks`` and ``manifest``; other ranks return empty structures but still
        participate in the distributed all-gathers.

        Returns:
            cpu_chunks:  List of ``{weight_name: cpu_tensor}`` dicts, one per shard file.
            manifest:    Mapping of ``weight_name → shard_filename``.
            total_size:  Total byte count of all exported tensors (for the index JSON).
        """
        cpu_chunks: list[dict[str, torch.Tensor]] = []
        manifest: dict[str, str] = {}
        current_chunk: dict[str, torch.Tensor] = {}
        current_chunk_bytes: int = 0
        total_size: int = 0
        file_idx: int = 0

        for name, param in model.model.model.named_parameters():
            # Phase 2+: HFModel initialises _model via AutoModelForImageTextToText /
            # AutoModelForCausalLM, so parameter names are HF-native and match the
            # safetensors checkpoint keys loaded by _load_vlm_weights().
            #
            # MoE note: Qwen3VLMoeTextExpertsGroupedMm stores expert weights in HF-native
            # grouped layout — gate_up_proj: [E, H, 2F], down_proj: [E, F, H] — matching
            # the checkpoint format exactly.  No transposition or per-expert fan-out is
            # needed.  (The legacy Phase 0 path stored tensors in a transposed internal
            # format [E, 2F, H] under the name "gate_and_up_projs" and required
            # weight_mapper.policy_map_local_key_for_export_tensor() to transpose back on
            # export.  Phase 2 uses HFModel and has no such internal reformat.)
            #
            # torch.compile and gradient-checkpointing wrappers inject prefixes into
            # named_parameters() output.  Strip them so exported keys are HF-native,
            # matching what HFModel._load_vlm_weights() does for the in-memory state dict.
            name = name.replace("_orig_mod.", "").replace("_checkpoint_wrapped_module.", "")

            # Gather across FSDP / TP / CP ranks (collective — all ranks must call).
            if isinstance(param, torch.distributed.tensor.DTensor):
                param = param.full_tensor()
            param = param.detach()
            if self._export_dtype is not None:
                param = param.to(dtype=self._export_dtype)

            tensor_bytes = param.element_size() * param.numel()

            # Flush the current chunk when the shard size limit would be exceeded.
            # current_chunk_bytes is tracked on ALL ranks so shard boundaries are
            # consistent (the shard_name written into manifest must agree everywhere).
            if current_chunk_bytes + tensor_bytes > self._MAX_SHARD_BYTES and current_chunk_bytes > 0:
                if is_rank0():
                    cpu_chunks.append(current_chunk)
                    current_chunk = {}
                file_idx += 1
                current_chunk_bytes = 0

            shard_name = f"{file_idx:05d}.safetensors"
            if is_rank0():
                current_chunk[name] = param.cpu()
                manifest[name] = shard_name
                total_size += tensor_bytes
            current_chunk_bytes += tensor_bytes

        # Flush the final (possibly partial) chunk.
        if current_chunk_bytes > 0 and is_rank0() and current_chunk:
            cpu_chunks.append(current_chunk)

        return cpu_chunks, manifest, total_size

    def _save_and_upload(
        self,
        cpu_chunks: list[dict[str, torch.Tensor]],
        manifest: dict[str, str],
        total_size: int,
        hf_config: Any,
        model_name_or_path: str,
        output_dir: str,
        iteration: int,
    ) -> None:
        """Write safetensors shards, HF config, tokenizer; upload to S3 / HF Hub.

        Runs on rank 0 inside a background thread.  Any exception is stored in
        ``self._worker_exception`` so the main thread can re-raise it.
        """
        try:
            self._do_save_and_upload(
                cpu_chunks, manifest, total_size, hf_config, model_name_or_path, output_dir, iteration
            )
        except Exception as exc:
            log.error(
                "[HFExportCallback] Export worker for iter %d raised an exception: %s",
                iteration,
                exc,
                exc_info=True,
            )
            self._worker_exception = exc

    def _do_save_and_upload(
        self,
        cpu_chunks: list[dict[str, torch.Tensor]],
        manifest: dict[str, str],
        total_size: int,
        hf_config: Any,
        model_name_or_path: str,
        output_dir: str,
        iteration: int,
    ) -> None:
        """Core export logic (called from the background thread via ``_save_and_upload``).

        Error handling is tiered:
        - Steps 1-4 (shards, index JSON, HF config, source-model file copy): any exception
          propagates to the outer ``_save_and_upload`` wrapper so the main thread is notified.
          A failed file copy leaves the checkpoint unusable for trust_remote_code models, so
          it is treated as a hard failure like the shard writes.
        - Steps 5-7 (tokenizer, generation_config, S3 upload, HF Hub upload): failures are
          treated as soft warnings.  The tokenizer and generation config are best-effort; upload
          failures do not invalidate the local safetensors export, so an outage must not abort
          training.
        """
        hf_cfg = self.config.checkpoint.hf_export
        os.makedirs(output_dir, exist_ok=True)
        log.info(f"[HFExportCallback] Writing iter {iteration} export to {output_dir}")

        # 1. Safetensors shards — one file per chunk (ordered by file_idx).
        #    Each chunk is cleared after writing so its tensors can be GC'd
        #    incrementally rather than being held until the whole loop completes.
        for i in range(len(cpu_chunks)):
            chunk = cpu_chunks[i]
            shard_path = os.path.join(output_dir, f"{i:05d}.safetensors")
            _safetensors_save_file(chunk, shard_path)
            log.info(f"[HFExportCallback] Wrote {shard_path}")
            cpu_chunks[i] = {}  # release tensor references for GC

        # 2. model.safetensors.index.json
        #    total_size is pre-computed in _gather_weights to avoid needing chunks here.
        index_json = {"metadata": {"total_size": total_size}, "weight_map": manifest}
        index_path = os.path.join(output_dir, "model.safetensors.index.json")
        with open(index_path, "w") as fh:
            json.dump(index_json, fh, indent=4)

        # 3. HuggingFace model config.
        hf_config.save_pretrained(output_dir)

        # 4. Copy missing .py/.json files for trust_remote_code models.
        #    Only applicable when model_name_or_path is a local directory.
        #    The full directory layout is preserved so nested packages referenced by
        #    auto_map are included (mirroring convert_checkpoint.py's copytree approach).
        #    Files already present in the export dir (e.g., config.json written by
        #    hf_config.save_pretrained) are never overwritten.
        #    HARD failure: a broken copy leaves the checkpoint unloadable, so any I/O error
        #    propagates to the background-worker wrapper (same as shard writes).
        if model_name_or_path and os.path.isdir(model_name_or_path):
            real_src = os.path.realpath(model_name_or_path)
            real_out = os.path.realpath(output_dir)
            copied = []
            for root, dirs, files in os.walk(real_src):
                real_root = os.path.realpath(root)
                # Prune any subtree that is, leads to, or is inside output_dir.
                # This prevents recursing into previously written export dirs when
                # output_dir (or a parent of it) lives inside model_name_or_path.
                dirs[:] = [
                    d
                    for d in dirs
                    if not (
                        (p := os.path.realpath(os.path.join(real_root, d))) == real_out
                        or p.startswith(real_out + os.sep)
                        or real_out.startswith(p + os.sep)
                    )
                ]
                if real_root == real_out or real_root.startswith(real_out + os.sep):
                    continue
                rel_dir = os.path.relpath(real_root, real_src)
                for fname in files:
                    if not (fname.endswith(".py") or fname.endswith(".json")):
                        continue
                    src = os.path.join(real_root, fname)
                    dst_dir = output_dir if rel_dir == "." else os.path.join(output_dir, rel_dir)
                    dst = os.path.join(dst_dir, fname)
                    if not os.path.exists(dst):
                        os.makedirs(dst_dir, exist_ok=True)
                        shutil.copy2(src, dst)
                        copied.append(os.path.join(rel_dir, fname) if rel_dir != "." else fname)
            if copied:
                log.info(f"[HFExportCallback] Copied missing files from source model: {copied}")

        # 5. Tokenizer (best-effort — may fail for custom / gated models).
        if AutoTokenizer is not None and model_name_or_path:
            try:
                tok = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True)
                tok.save_pretrained(output_dir)
            except Exception as exc:
                log.warning(f"[HFExportCallback] Tokenizer save skipped: {exc}")

        # 6. Generation config (best-effort — not all models expose one).
        if GenerationConfig is not None and model_name_or_path:
            try:
                gen_cfg = GenerationConfig.from_pretrained(model_name_or_path, trust_remote_code=True)
                gen_cfg.save_pretrained(output_dir)
            except Exception as exc:
                log.warning(f"[HFExportCallback] generation_config save skipped: {exc}")

        # 7. S3 upload — soft failure: local export is intact regardless of upload outcome.
        obj_store = hf_cfg.upload_to_object_store
        if obj_store.enabled:
            s3_prefix = f"{self.config.job.path}/hf_exports/iter_{iteration:09d}"
            try:
                _upload_folder_to_s3(output_dir, obj_store.bucket, s3_prefix, obj_store.credentials)
                log.info(f"[HFExportCallback] S3 upload done: s3://{obj_store.bucket}/{s3_prefix}")
            except Exception as exc:
                # Intentionally soft: an upload outage must not crash training.
                log.warning(f"[HFExportCallback] S3 upload failed (local export intact): {exc}")

        # 8. HuggingFace Hub upload — soft failure: see comment above.
        if hf_cfg.hf_repo_id:
            self._upload_to_hf_hub(output_dir, hf_cfg.hf_repo_id)

        log.info(f"[HFExportCallback] Export complete for iter {iteration}.")

    @staticmethod
    def _upload_to_hf_hub(output_dir: str, repo_id: str, max_retries: int = 3) -> None:
        try:
            from huggingface_hub import HfApi
        except ImportError:
            log.warning("[HFExportCallback] huggingface_hub not installed; skipping HF Hub upload.")
            return

        api = HfApi()
        for attempt in range(1, max_retries + 1):
            try:
                api.create_repo(repo_id=repo_id, exist_ok=True)
                break
            except Exception as exc:
                log.warning(f"[HFExportCallback] create_repo attempt {attempt}/{max_retries} failed: {exc}")
                if attempt == max_retries:
                    log.warning(
                        f"[HFExportCallback] Could not create HF Hub repo '{repo_id}' after "
                        f"{max_retries} attempts; skipping upload."
                    )
                    return

        for attempt in range(1, max_retries + 1):
            try:
                api.upload_folder(
                    folder_path=output_dir,
                    repo_id=repo_id,
                    commit_message=f"Upload checkpoint from {os.path.basename(output_dir)}",
                )
                log.info(f"[HFExportCallback] Uploaded to HF Hub: {repo_id}")
                return
            except Exception as exc:
                log.warning(f"[HFExportCallback] HF Hub upload attempt {attempt}/{max_retries} failed: {exc}")

        log.warning(f"[HFExportCallback] All {max_retries} HF Hub upload attempts failed for {repo_id}.")