Cache examples and simplify frontend

.gitattributes

@@ -36,3 +36,32 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 examples/bowery_waltz.mp4 filter=lfs diff=lfs merge=lfs -text
 examples/coyote_run.mp4 filter=lfs diff=lfs merge=lfs -text
 examples/shove_it.mp4 filter=lfs diff=lfs merge=lfs -text
+candidate_examples/brick_texture.jpg filter=lfs diff=lfs merge=lfs -text
+candidate_examples/contact_sheet.jpg filter=lfs diff=lfs merge=lfs -text
+candidate_examples/dragonfly_macro.jpg filter=lfs diff=lfs merge=lfs -text
+candidate_examples/embroidery_threads.jpg filter=lfs diff=lfs merge=lfs -text
+candidate_examples/feather_barbs.jpg filter=lfs diff=lfs merge=lfs -text
+candidate_examples/leaf_macro.jpg filter=lfs diff=lfs merge=lfs -text
+candidate_examples/palm_bark.jpg filter=lfs diff=lfs merge=lfs -text
+candidate_results/brick_texture_crop_compare.jpg filter=lfs diff=lfs merge=lfs -text
+candidate_results/brick_texture_pid_preview.jpg filter=lfs diff=lfs merge=lfs -text
+candidate_results/dragonfly_macro_crop_compare.jpg filter=lfs diff=lfs merge=lfs -text
+candidate_results/dragonfly_macro_pid_preview.jpg filter=lfs diff=lfs merge=lfs -text
+candidate_results/embroidery_threads_crop_compare.jpg filter=lfs diff=lfs merge=lfs -text
+candidate_results/embroidery_threads_pid_preview.jpg filter=lfs diff=lfs merge=lfs -text
+candidate_results/feather_barbs_crop_compare.jpg filter=lfs diff=lfs merge=lfs -text
+candidate_results/feather_barbs_pid_preview.jpg filter=lfs diff=lfs merge=lfs -text
+candidate_results/leaf_macro_crop_compare.jpg filter=lfs diff=lfs merge=lfs -text
+candidate_results/leaf_macro_pid_preview.jpg filter=lfs diff=lfs merge=lfs -text
+candidate_results/palm_bark_crop_compare.jpg filter=lfs diff=lfs merge=lfs -text
+candidate_results/palm_bark_pid_preview.jpg filter=lfs diff=lfs merge=lfs -text
+examples/brick_wall_texture.jpg filter=lfs diff=lfs merge=lfs -text
+examples/leaf_macro_texture.jpg filter=lfs diff=lfs merge=lfs -text
+examples/palm_bark_texture.jpg filter=lfs diff=lfs merge=lfs -text
+research_outputs/alpine_absdiff_x10.png filter=lfs diff=lfs merge=lfs -text
+research_outputs/alpine_absdiff_x8.png filter=lfs diff=lfs merge=lfs -text
+research_outputs/alpine_cabin_compare.png filter=lfs diff=lfs merge=lfs -text
+research_outputs/alpine_foreground_compare.png filter=lfs diff=lfs merge=lfs -text
+research_outputs/alpine_lanczos_to_4k.png filter=lfs diff=lfs merge=lfs -text
+research_outputs/alpine_peaks_compare.png filter=lfs diff=lfs merge=lfs -text
+research_outputs/alpine_pid_4k.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,5 @@
+/.cache/
+/checkpoints/
+__pycache__/
+*.pyc
+.DS_Store

LICENSE-PiD-CODE

@@ -0,0 +1,13 @@
+Copyright 2026 NVIDIA CORPORATION & AFFILIATES.
+
+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.

UPSTREAM_PID_README.md

@@ -0,0 +1,245 @@
+# PiD — Pixel Diffusion Decoder
+
+> **TL;DR** — PiD is a plug-and-play diffusion decoder that replaces VAE/RAE decoders, turning latent representations directly into super-resolved pixels in a single pass.
+
+<p align="center">
+  <img src="figures/teaser.jpg" alt="PiD teaser" width="100%">
+</p>
+
+https://github.com/user-attachments/assets/a556e2d4-5de5-4bcf-9daa-80f7ea6b2124
+
+PiD reformulates the latent-to-pixel decoder as a conditional pixel-space diffusion
+model, unifying decoding and upsampling into a single generative module.
+It directly denoises in high-resolution pixel
+space and produces a super-resolved image in one pass.
+
+**[Paper](https://arxiv.org/abs/2605.23902), [Project Page](https://research.nvidia.com/labs/sil/projects/pid/), [Model Weights](https://huggingface.co/nvidia/PiD)**
+
+[Yifan Lu](https://yifanlu0227.github.io/),
+[Qi Wu](https://wilsoncernwq.github.io/),
+[Jay Zhangjie Wu](https://zhangjiewu.github.io/),
+[Zian Wang](https://www.cs.toronto.edu/~zianwang/),
+[Huan Ling](https://www.cs.toronto.edu/~linghuan/),
+[Sanja Fidler](https://www.cs.utoronto.ca/~fidler/),
+[Xuanchi Ren](https://xuanchiren.com/) <br>
+
+## Installation
+
+> [!TIP]
+> **Quick Start** — if your environment already has PyTorch (with CUDA), `transformers>=4.57.x`, and `diffusers>=0.37`, you don't need to build a new conda env. Just install the small set of utility deps the inference code pulls eagerly and you're ready to run the diffusers backbones (`flux`/`flux2`/`sd3`/`zimage`):
+>
+> ```bash
+> pip install hydra-core==1.3.2 omegaconf==2.3.0 \
+>     attrs einops loguru termcolor fvcore iopath pynvml wandb \
+>     imageio opencv-python-headless pandas \
+>     safetensors "huggingface-hub>=1.0" sentencepiece boto3 botocore
+> pip install -e .
+> ```
+>
+> For the `dinov2` / `siglip` backbones you additionally need the upstream RAE / Scale-RAE repos plus a couple of extra packages — see [docs/dinov2_siglip.md](docs/dinov2_siglip.md).
+
+Full conda-managed install (preferred if you're starting from scratch):
+
+```bash
+conda env create -f environment.yml
+conda activate pid
+
+# 2. Install this package in editable mode.
+pip install -e .
+```
+
+## Checkpoints and assets
+
+Pretrained PiD checkpoints live under `checkpoints/`. Each diffusers backbone ships
+two variants — the original `2k` decoder (trained at 2048px) and a `2kto4k` decoder
+(trained with multi-resolution data bucketing 2048→3840 + an SD3-style dynamic
+shift, intended for 1024 LDM → 4K decoding). Pick the variant at the CLI via
+`--pid_ckpt_type {2k,2kto4k}` (default: `2k`).
+
+### Downloading
+
+The released decoder weights and the encoder/decoder ("VAE") weights they
+depend on are hosted at [`nvidia/PiD`](https://huggingface.co/nvidia/PiD) on
+the Hugging Face Hub. Pull just the `checkpoints/` tree into this repo:
+
+```bash
+hf download nvidia/PiD --local-dir . --include "checkpoints/*"
+```
+
+## Running inference
+
+PiD ships two complementary entry points per backbone:
+
+| Backbone | `from_clean_*` (image → encode → PiD) | `from_ldm_*` (text/class → LDM → PiD) |
+|----------|---------------------------------------|---------------------------------------|
+| flux     | `from_clean_flux.py`    | `from_ldm_flux.py`    |
+| flux2    | `from_clean_flux2.py`   | `from_ldm_flux2.py`   |
+| sd3      | `from_clean_sd3.py`     | `from_ldm_sd3.py`     |
+| zimage   | reuses `flux`           | `from_ldm_zimage.py`  |
+| dinov2   | `from_clean_dinov2.py`  | `from_ldm_dinov2.py`  |
+| siglip   | `from_clean_siglip.py`  | `from_ldm_siglip.py`  |
+
+All scripts live under `pid/_src/inference/` and decode each captured latent
+twice — once with the backbone's native VAE (baseline) and once with PiD.
+
+> [!IMPORTANT]
+> Picking the checkpoint variant — `--pid_ckpt_type`
+> Every entry point accepts `--pid_ckpt_type {2k,2kto4k}` (default `2k`):
+>
+> - **`2k`** — the original 2048px-trained decoder.
+> - **`2kto4k`** — the up-to-4K-resolution decoder. > > Available for `flux` / `flux2` / `sd3` / `zimage` only. Worse than `2k` at 2048px resolution.
+>
+> For the exact checkpoint path for each backbone, see [docs/checkpoints.md](docs/checkpoints.md).
+> A quick sanity check that the right variant loaded: when `2kto4k` is active you
+should see `PixelDiT dynamic shift: base_shift=4.0 base_image_size=1024` in the
+init log; for `2k` that line is absent. Both `2k` and `2kto4k` support non-square aspect ratios.
+
+### 📕 `from_ldm_*`: text / class → latent diffusion → PiD decode
+
+Runs the corresponding latent-diffusion backbone on a prompt (or class id for
+the class-conditional `dinov2` backbone), captures the intermediate `x_t` at
+user-specified denoising steps (early LDM termination) and the final clean `x_0`, then decodes
+each captured latent with both the native VAE / RAE decoder (baseline) and PiD.
+
+For `flux` / `flux2` / `sd3` / `zimage` the LDM is a HuggingFace `diffusers`
+pipeline (`FluxPipeline`, `Flux2Pipeline`, `StableDiffusion3Pipeline`,
+`ZImagePipeline`).
+
+For `dinov2` and `siglip` the LDM is the upstream
+[RAE](https://github.com/bytetriper/RAE) (class-conditional ImageNet-512) or
+[Scale-RAE](https://github.com/ZitengWangNYU/Scale-RAE) (text-conditional
+256px) repo — see the optional-deps section below for installation.
+
+#### Example 1 — Single-GPU, single prompt (Flux, default `2k` decoder)
+
+```bash
+PYTHONPATH=. python -m pid._src.inference.from_ldm_flux \
+    --prompt "A photorealistic half-body portrait of a brown tabby cat with bold stripes sitting attentively on a rustic wooden kitchen table, soft morning light streaming sideways through a large window, fine fur detail and stripe patterns sharply visible, intense amber-green eyes in razor-sharp focus, warm farmhouse kitchen softly out of focus, cinematic shallow depth of field, ultra-detailed fur texture, photorealistic" \
+    --ldm_inference_steps 28 --save_xt_steps 24 \
+    --output_dir ./results/official_demo/flux \
+    --cfg_scale 1 --pid_inference_steps 4 --scale 4
+```
+
+#### Example 2 — Single-GPU, 4K decode (Flux, `2kto4k` decoder)
+
+Same backbone as Example 1 but with `--resolution 1024 --pid_ckpt_type 2kto4k`,
+so the LDM produces a 1024² latent and PiD decodes it to 4K.
+
+```bash
+PYTHONPATH=. python -m pid._src.inference.from_ldm_flux \
+    --prompt "A photorealistic half-body portrait of a brown tabby cat with bold stripes sitting attentively on a rustic wooden kitchen table, soft morning light streaming sideways through a large window, fine fur detail and stripe patterns sharply visible, intense amber-green eyes in razor-sharp focus, warm farmhouse kitchen softly out of focus, cinematic shallow depth of field, ultra-detailed fur texture, photorealistic" \
+    --resolution 1024 --pid_ckpt_type 2kto4k \
+    --ldm_inference_steps 28 --save_xt_steps 24 \
+    --output_dir ./results/official_demo/flux_4k \
+    --cfg_scale 1 --pid_inference_steps 4 --scale 4
+```
+
+#### Example 3 — Multi-GPU with a prompt file (Z-Image)
+
+`torchrun` shards `--prompt_file` across ranks; each rank writes to
+`--output_dir` independently.
+
+```bash
+PYTHONPATH=. torchrun --nproc_per_node=4 \
+    -m pid._src.inference.from_ldm_zimage \
+    --prompt_file pid/_src/inference/prompts/prompt_creative.txt \
+    --ldm_inference_steps 50 --save_xt_steps 46 \
+    --output_dir ./results/official_demo/zimage \
+    --cfg_scale 1 --pid_inference_steps 4 --scale 4
+```
+
+#### `dinov2` / `siglip` backbones
+
+The upstream RAE / Scale-RAE LDMs don't live in `diffusers` — see
+[`docs/dinov2_siglip.md`](docs/dinov2_siglip.md) for setup and end-to-end
+examples.
+
+#### Suggested step settings per diffusers backbone
+
+(See each script's docstring for the exact recipe.)
+
+| Backbone | LDM steps flag          | Default steps | `--save_xt_steps` (example) | Best `--save_xt_steps` |
+|----------|-------------------------|---------------|-----------------------------|----------------------|
+| flux     | `--ldm_inference_steps` | 28            | `22 24 26`         | 24  |
+| sd3      | `--ldm_inference_steps` | 28            | `22 24 26`         | 24  |
+| flux2    | `--ldm_inference_steps` | 50            | `44 46 48`         | 46  |
+| zimage   | `--ldm_inference_steps` | 50            | `44 46 48`         | 46  |
+
+---
+### 📗 `from_clean_*`: image → VAE encode → PiD decode
+
+No latent diffusion model is run. The input image is encode by VAE,
+optionally corrupted with Gaussian noise at each
+sigma in `--degrade_sigmas`, then decoded by PiD at `--scale * input_resolution`.
+
+Single-GPU example (Flux):
+
+```bash
+PYTHONPATH=. python -m pid._src.inference.from_clean_flux \
+    --manifest assets/clean_image_manifest.jsonl \
+    --input_resolution 512 \
+    --degrade_sigmas 0.0 \
+    --output_dir ./results/official_demo_from_clean/flux \
+    --cfg_scale 1 --pid_inference_steps 4 --scale 4
+```
+
+You can pass a single image with `--input_path` and a prompt with `--prompt`
+instead of `--manifest`, and a sigma sweep such as `--degrade_sigmas 0.0 0.2 0.4 0.8`
+to decode noise-corrupted latents.
+
+The `dinov2` / `siglip` `from_clean_*` flows take the same flags but with
+different default resolutions and scales —
+see [`docs/dinov2_siglip.md`](docs/dinov2_siglip.md).
+
+### Common arguments
+
+| Flag | Meaning |
+|------|---------|
+| `--pid_inference_steps`| Number of denoising steps for PiD (4 for the released distilled checkpoints) |
+| `--scale`              | PiD upscale factor (output = `baseline * scale`); 8 for Scale-RAE and 4 for other backbones |
+| `--cfg_scale`          | Classifier-free guidance scale for PiD |
+| `--output_dir`         | Where to write the side-by-side comparison images |
+| `--seed`               | Base random seed |
+
+Multi-GPU runs use `torchrun --nproc_per_node=N`; each rank processes a shard
+of the prompts / manifest entries and writes to `--output_dir` independently.
+
+## Repository layout
+
+```
+pid/_src/inference/
+├── from_ldm_{flux,flux2,sd3,zimage,dinov2,siglip}.py  # text/class → LDM → PiD decode
+├── from_clean_{flux,flux2,sd3,dinov2,siglip}.py       # image → encode → PiD decode
+├── _demo_common.py                                    # shared CLI + run loop for from_ldm_*
+├── _demo_from_clean_common.py                         # shared CLI + run loop for from_clean_*
+├── checkpoint_registry.py                             # backbone → PiD checkpoint mapping
+├── pipeline_registry.py                               # diffusers backbone → HF pipeline mapping
+├── rae_generation.py                                  # DINOv2-RAE LDM helpers (from_ldm_dinov2)
+├── scale_rae_generation.py                            # Scale-RAE LDM helpers (from_ldm_siglip)
+└── prompts/                                           # prompt files for from_ldm_*
+```
+
+## License
+
+PiD codebase is licensed under the [Apache License 2.0](LICENSE).
+
+## Contributing
+
+See [`CONTRIBUTING.md`](CONTRIBUTING.md) for development setup, code style,
+and the DCO sign-off requirement.
+
+## Acknowledgments
+
+The authors would like to acknowledge [Yongsheng Yu](https://www.yongshengyu.com/) and [Wei Xiong](https://wxiong.me/) for open-sourcing [PixelDiT](https://pixeldit.github.io/)'s model and weights, and thank Product Managers [Aditya Mahajan](https://www.linkedin.com/in/aditya-mahajan1) and [Matt Cragun](https://www.linkedin.com/in/mcragun/) for their valuable support and guidance.
+
+
+## Citation
+
+```bibtex
+@article{lu2026pid,
+    title={PiD: Fast and High-Resolution Latent Decoding with Pixel Diffusion},
+    author={Lu, Yifan and Wu, Qi and Wu, Jay Zhangjie and Wang, Zian and Ling, Huan and Fidler, Sanja and Ren, Xuanchi},
+    journal={arXiv preprint arXiv:2605.23902},
+    year={2026}
+}
+```

app.py

@@ -40,6 +40,7 @@ from gradio_client.data_classes import FileData
 APP_ROOT = Path(__file__).resolve().parent
 INDEX_PATH = APP_ROOT / "index.html"
 EXAMPLES_DIR = APP_ROOT / "examples"
+EXAMPLE_STATIC_CACHE_DIR = APP_ROOT / "example_cache"
 EXAMPLE_CACHE_DIR = Path("/tmp/marlin_example_cache")
 MODEL_ID = "NemoStation/Marlin-2B"
 MODEL_DESCRIPTION = (
@@ -241,18 +242,23 @@ def _zerogpu_probe() -> str:
     return "ready"
 
 
+def _static_example_cache_path(index: int) -> Path:
+    return EXAMPLE_STATIC_CACHE_DIR / f"v2_example_{index}.json"
+
+
 def _example_cache_path(index: int) -> Path:
     return EXAMPLE_CACHE_DIR / f"v2_example_{index}.json"
 
 
 def _read_example_cache(index: int) -> dict[str, Any] | None:
-    path = _example_cache_path(index)
-    if not path.exists():
-        return None
-    try:
-        return json.loads(path.read_text(encoding="utf-8"))
-    except Exception:
-        return None
+    for path in (_static_example_cache_path(index), _example_cache_path(index)):
+        if not path.exists():
+            continue
+        try:
+            return json.loads(path.read_text(encoding="utf-8"))
+        except Exception:
+            continue
+    return None
 
 
 def _write_example_cache(index: int, payload: dict[str, Any]) -> None:
@@ -282,14 +288,21 @@ async def health() -> JSONResponse:
 
 @app.get("/examples.json")
 async def examples_json() -> JSONResponse:
+    def _example_payload(index: int, example: dict[str, Any]) -> dict[str, Any]:
+        thumbnail = Path(example["file"]).with_suffix(".jpg").name
+        payload = {
+            **example,
+            "url": f"/examples/{example['file']}",
+            "cached": _read_example_cache(index) is not None,
+        }
+        if (EXAMPLES_DIR / thumbnail).exists():
+            payload["thumbnail"] = f"/examples/{thumbnail}"
+        return payload
+
     return JSONResponse(
         {
             "examples": [
-                {
-                    **example,
-                    "url": f"/examples/{example['file']}",
-                    "cached": _read_example_cache(i) is not None,
-                }
+                _example_payload(i, example)
                 for i, example in enumerate(EXAMPLES)
             ]
         }

candidate_results/summary.json

@@ -0,0 +1,104 @@
+[
+  {
+    "slug": "brick_texture",
+    "status": "14.63s GPU compute | 1024x1024 input -> 4096x4096 output | 2kto4k | slider baseline: Lanczos input resize",
+    "mean_abs": 12.956155200799307,
+    "rms": 19.173621112407616,
+    "max_delta": 226,
+    "pct_gt_2": 94.4489061832428,
+    "pct_gt_5": 69.04728412628174,
+    "pct_gt_10": 41.335028409957886,
+    "crop_box": [
+      2176,
+      2368,
+      2944,
+      3136
+    ],
+    "prompt": "A close-up photograph of a weathered brick wall with rough mortar lines, chipped corners, sand grains, and subtle orange-brown surface texture."
+  },
+  {
+    "slug": "dragonfly_macro",
+    "status": "14.92s GPU compute | 1024x1024 input -> 4096x4096 output | 2kto4k | slider baseline: Lanczos input resize",
+    "mean_abs": 4.303538699944814,
+    "rms": 5.7793400659609215,
+    "max_delta": 133,
+    "pct_gt_2": 81.09467029571533,
+    "pct_gt_5": 24.35404062271118,
+    "pct_gt_10": 4.072648286819458,
+    "crop_box": [
+      1088,
+      1216,
+      1856,
+      1984
+    ],
+    "prompt": "A macro photograph of a dragonfly with glassy compound eyes, translucent veined wings, tiny hairs, and natural outdoor light."
+  },
+  {
+    "slug": "embroidery_threads",
+    "status": "15.01s GPU compute | 1024x1024 input -> 4096x4096 output | 2kto4k | slider baseline: Lanczos input resize",
+    "mean_abs": 4.275510291258494,
+    "rms": 5.43537963982756,
+    "max_delta": 80,
+    "pct_gt_2": 87.73354887962341,
+    "pct_gt_5": 27.23088264465332,
+    "pct_gt_10": 1.5419423580169678,
+    "crop_box": [
+      2560,
+      3328,
+      3328,
+      4096
+    ],
+    "prompt": "A close-up photograph of embroidered fabric with raised yellow and green thread, visible woven cloth texture, stitch fibers, and curved ornamental patterns."
+  },
+  {
+    "slug": "feather_barbs",
+    "status": "15.08s GPU compute | 1024x1024 input -> 4096x4096 output | 2kto4k | slider baseline: Lanczos input resize",
+    "mean_abs": 3.7417769034703574,
+    "rms": 4.578883124584203,
+    "max_delta": 84,
+    "pct_gt_2": 76.96071863174438,
+    "pct_gt_5": 20.45971155166626,
+    "pct_gt_10": 1.09291672706604,
+    "crop_box": [
+      2112,
+      3328,
+      2880,
+      4096
+    ],
+    "prompt": "A close-up photograph of soft white feathers with fine barbs, layered strands, subtle shadows, and delicate down texture."
+  },
+  {
+    "slug": "leaf_macro",
+    "status": "14.62s GPU compute | 1024x1024 input -> 4096x4096 output | 2kto4k | slider baseline: Lanczos input resize",
+    "mean_abs": 11.487248241901398,
+    "rms": 16.521660198582843,
+    "max_delta": 182,
+    "pct_gt_2": 97.48759865760803,
+    "pct_gt_5": 74.37732815742493,
+    "pct_gt_10": 37.26730942726135,
+    "crop_box": [
+      0,
+      1216,
+      768,
+      1984
+    ],
+    "prompt": "A macro photograph of a green leaf surface with branching veins, tiny cells, crisp organic texture, and backlit detail."
+  },
+  {
+    "slug": "palm_bark",
+    "status": "17.11s GPU compute | 1024x1024 input -> 4096x4096 output | 2kto4k | slider baseline: Lanczos input resize",
+    "mean_abs": 27.973096787929535,
+    "rms": 36.79712697730481,
+    "max_delta": 242,
+    "pct_gt_2": 98.77586364746094,
+    "pct_gt_5": 90.33623933792114,
+    "pct_gt_10": 74.35104846954346,
+    "crop_box": [
+      2496,
+      1664,
+      3264,
+      2432
+    ],
+    "prompt": "A close-up photograph of palm tree bark with deep vertical cracks, rough fibers, beige-gray ridges, and natural wood texture."
+  }
+]

example_cache/v2_example_0.json

@@ -0,0 +1,44 @@
+{
+  "ok": true,
+  "task": "Caption",
+  "scene": "The scene is set on a minimalist stage with a stark black background, creating a high-contrast environment that emphasizes the performers. A man and a woman are the central figures, dressed in mid-20th-century attire. The man wears a light-colored, long-sleeved jacket and matching trousers, while the woman is in a light-colored, knee-length dress with a dark belt and a small hat. The stage floor is a flat, light-colored surface that reflects the overhead lighting. The camera maintains a medium-wide shot, capturing the full bodies of both dancers as they move across the stage. The atmosphere is theatrical and energetic, characteristic of a choreographed dance performance. The lighting is focused primarily on the couple, leaving the edges of the stage in deep shadow. The man's jacket is slightly oversized, and the woman's dress has a full skirt that flares out during her movements. The overall aesthetic is that of a classic black-and-white film, with visible film grain and occasional flickers.",
+  "events": [
+    {
+      "start": 0.5,
+      "end": 1.5,
+      "description": "The man and woman walk toward each other."
+    },
+    {
+      "start": 1.5,
+      "end": 2.5,
+      "description": "The woman bends down to touch the man's hand."
+    },
+    {
+      "start": 2.5,
+      "end": 4.5,
+      "description": "The man and woman begin dancing together."
+    },
+    {
+      "start": 4.5,
+      "end": 6.5,
+      "description": "The man spins the woman in a circle."
+    },
+    {
+      "start": 6.5,
+      "end": 10.5,
+      "description": "The couple performs synchronized side-to-side steps."
+    },
+    {
+      "start": 10.5,
+      "end": 14.0,
+      "description": "The man leads the woman through a final spin."
+    }
+  ],
+  "span": null,
+  "raw": "Scene: The scene is set on a minimalist stage with a stark black background, creating a high-contrast environment that emphasizes the performers. A man and a woman are the central figures, dressed in mid-20th-century attire. The man wears a light-colored, long-sleeved jacket and matching trousers, while the woman is in a light-colored, knee-length dress with a dark belt and a small hat. The stage floor is a flat, light-colored surface that reflects the overhead lighting. The camera maintains a medium-wide shot, capturing the full bodies of both dancers as they move across the stage. The atmosphere is theatrical and energetic, characteristic of a choreographed dance performance. The lighting is focused primarily on the couple, leaving the edges of the stage in deep shadow. The man's jacket is slightly oversized, and the woman's dress has a full skirt that flares out during her movements. The overall aesthetic is that of a classic black-and-white film, with visible film grain and occasional flickers.\n\nEvents:\n<0.5 - 1.5> The man and woman walk toward each other.\n<1.5 - 2.5> The woman bends down to touch the man's hand.\n<2.5 - 4.5> The man and woman begin dancing together.\n<4.5 - 6.5> The man spins the woman in a circle.\n<6.5 - 10.5> The couple performs synchronized side-to-side steps.\n<10.5 - 14.0> The man leads the woman through a final spin.",
+  "status": "Caption completed in 9.1s on GPU | 14.0s video | 6 events parsed",
+  "duration": 14.0,
+  "timings": {
+    "gpu_seconds": 9.12043823604472
+  }
+}

example_cache/v2_example_1.json

@@ -0,0 +1,23 @@
+{
+  "ok": true,
+  "task": "Find",
+  "scene": "",
+  "events": [
+    {
+      "start": 0.0,
+      "end": 1.5,
+      "description": "the skateboard spins under the rider's feet"
+    }
+  ],
+  "span": {
+    "start": 0.0,
+    "end": 1.5,
+    "label": "the skateboard spins under the rider's feet"
+  },
+  "raw": "From 0.0 to 1.5.",
+  "status": "Find completed in 0.5s on GPU | 6.5s video | format_ok=True",
+  "duration": 6.5,
+  "timings": {
+    "gpu_seconds": 0.5369779840111732
+  }
+}

example_cache/v2_example_2.json

@@ -0,0 +1,23 @@
+{
+  "ok": true,
+  "task": "Find",
+  "scene": "",
+  "events": [
+    {
+      "start": 0.0,
+      "end": 1.5,
+      "description": "the coyote runs across the scene"
+    }
+  ],
+  "span": {
+    "start": 0.0,
+    "end": 1.5,
+    "label": "the coyote runs across the scene"
+  },
+  "raw": "From 0.0 to 1.5.",
+  "status": "Find completed in 0.5s on GPU | 9.9s video | format_ok=True",
+  "duration": 9.916667,
+  "timings": {
+    "gpu_seconds": 0.5370825119316578
+  }
+}

index.html

@@ -28,6 +28,7 @@
       --font: "Inter", system-ui, -apple-system, "SF Pro Text", sans-serif;
     }
     * { box-sizing: border-box; }
+    [hidden] { display: none !important; }
     html, body {
       margin: 0;
       min-height: 100%;
@@ -68,7 +69,11 @@
     .layout {
       display: grid;
       grid-template-columns: 330px minmax(0, 1fr);
+      width: min(calc(100% - 32px), 1380px);
       min-height: calc(100vh - 44px);
+      margin: 0 auto;
+      border-left: 1px solid var(--border);
+      border-right: 1px solid var(--border);
     }
     .column { min-width: 0; min-height: 0; border-right: 1px solid var(--border); }
     .viewer-column { border-right: 0; display: grid; grid-template-rows: auto minmax(0, 1fr); }
@@ -123,16 +128,22 @@
     .example-list { display: grid; gap: 6px; }
     .example-btn {
       text-align: left; border: 1px solid var(--border); border-radius: var(--radius);
-      background: var(--surface); padding: 9px 10px; display: grid; gap: 2px;
+      background: var(--surface); padding: 8px; display: grid;
+      grid-template-columns: 68px minmax(0, 1fr); gap: 10px; align-items: center;
     }
     .example-btn:hover, .example-btn.active { background: var(--surface-2); border-color: var(--border-strong); }
-    .example-btn strong { font-size: 12.5px; font-weight: 500; }
-    .example-btn span { color: var(--muted); font-size: 11.5px; }
+    .example-thumb {
+      width: 68px; height: 42px; object-fit: cover; border-radius: 4px;
+      background: #050506; border: 1px solid var(--border);
+    }
+    .example-copy { min-width: 0; display: grid; gap: 2px; }
+    .example-copy strong { font-size: 12.5px; font-weight: 500; white-space: nowrap; overflow: hidden; text-overflow: ellipsis; }
+    .example-copy span { color: var(--muted); font-size: 11.5px; }
     .viewer-header { display: flex; align-items: center; justify-content: space-between; gap: 10px; }
     .status { display: inline-flex; align-items: center; gap: 7px; color: var(--muted); font-size: 12px; }
     .status svg { color: var(--muted-2); }
-    .viewer { padding: 18px; min-height: 0; display: grid; grid-template-columns: minmax(0, 1.1fr) minmax(300px, .9fr); gap: 18px; }
-    .stage, .results {
+    .viewer { padding: 18px; min-height: 0; display: block; }
+    .stage {
       min-width: 0; border: 1px solid var(--border); border-radius: var(--radius);
       background: var(--surface); overflow: hidden;
     }
@@ -155,7 +166,7 @@
     }
     .segment.active { background: rgba(244,244,245,.72); }
     .time-readout { display: flex; justify-content: space-between; color: var(--muted); font-size: 11.5px; }
-    .events { display: grid; gap: 6px; max-height: 320px; overflow: auto; padding-right: 2px; }
+    .events { display: grid; gap: 6px; padding-right: 2px; }
     .event-row {
       width: 100%; text-align: left; border: 1px solid var(--border); border-radius: var(--radius);
       background: var(--surface); padding: 9px 10px; display: grid; gap: 3px;
@@ -163,7 +174,6 @@
     .event-row:hover, .event-row.active { background: var(--surface-2); border-color: var(--border-strong); }
     .event-row time { color: var(--muted); font-size: 11.5px; }
     .event-row span { color: var(--ink); }
-    .results { padding: 14px; display: grid; grid-template-rows: auto auto minmax(0,1fr); gap: 14px; }
     .section { display: grid; gap: 7px; min-width: 0; }
     .section h3 { margin: 0; color: var(--muted); font-size: 11px; text-transform: uppercase; letter-spacing: .05em; font-weight: 500; }
     .scene, .raw {
@@ -172,6 +182,28 @@
       padding: 10px; overflow: auto;
     }
     .raw { color: var(--muted); max-height: 260px; }
+    .model-details {
+      border: 1px solid var(--border); border-radius: var(--radius);
+      background: var(--surface); overflow: hidden;
+    }
+    .model-details summary {
+      min-height: 38px; padding: 0 10px; display: flex; align-items: center; justify-content: space-between; gap: 12px;
+      cursor: pointer; color: var(--muted); list-style: none;
+    }
+    .model-details summary::-webkit-details-marker { display: none; }
+    .model-details summary span:first-child {
+      font-size: 11px; font-weight: 500; text-transform: uppercase; letter-spacing: .05em;
+    }
+    .model-details summary span:last-child {
+      min-width: 0; overflow: hidden; white-space: nowrap; text-overflow: ellipsis; font-size: 11.5px;
+    }
+    .model-details summary::after {
+      content: ""; width: 7px; height: 7px; border-right: 1px solid var(--muted); border-bottom: 1px solid var(--muted);
+      transform: rotate(45deg); flex: 0 0 auto; transition: transform .15s ease;
+    }
+    .model-details[open] summary { border-bottom: 1px solid var(--border); }
+    .model-details[open] summary::after { transform: rotate(225deg); }
+    .details-body { padding: 12px; display: grid; gap: 14px; }
     .empty { color: var(--muted-2); }
     .toast {
       position: fixed; left: 50%; bottom: 18px; transform: translateX(-50%) translateY(14px);
@@ -181,9 +213,8 @@
     }
     .toast.show { opacity: 1; transform: translateX(-50%) translateY(0); }
     @media (max-width: 1080px) {
-      .layout { grid-template-columns: 1fr; }
+      .layout { grid-template-columns: 1fr; width: 100%; border-left: 0; border-right: 0; }
       .column { border-right: 0; border-bottom: 1px solid var(--border); }
-      .viewer { grid-template-columns: 1fr; }
     }
     @media (max-width: 680px) {
       .topbar { padding: 0 12px; }
@@ -274,23 +305,25 @@
               <div class="timebar" id="timebar"><div class="progress-line" id="progressLine"></div></div>
               <div class="time-readout"><span id="currentTime">0.0s</span><span id="durationTime">0.0s</span></div>
               <div class="events" id="eventsList"></div>
+              <details class="model-details" id="modelDetails">
+                <summary><span>Model output</span><span id="detailsSummary">Idle</span></summary>
+                <div class="details-body">
+                  <div class="section">
+                    <h3>Scene</h3>
+                    <p class="scene empty" id="sceneText">No scene yet.</p>
+                  </div>
+                  <div class="section">
+                    <h3>Status</h3>
+                    <p class="scene" id="runStatus">Idle</p>
+                  </div>
+                  <div class="section">
+                    <h3>Raw output</h3>
+                    <pre class="raw" id="rawText"></pre>
+                  </div>
+                </div>
+              </details>
             </div>
           </section>
-
-          <aside class="results">
-            <div class="section">
-              <h3>Scene</h3>
-              <p class="scene empty" id="sceneText">No scene yet.</p>
-            </div>
-            <div class="section">
-              <h3>Status</h3>
-              <p class="scene" id="runStatus">Idle</p>
-            </div>
-            <div class="section">
-              <h3>Raw output</h3>
-              <pre class="raw" id="rawText"></pre>
-            </div>
-          </aside>
         </div>
       </section>
     </section>
@@ -337,9 +370,19 @@
       sceneText: $("sceneText"),
       rawText: $("rawText"),
       runStatus: $("runStatus"),
+      detailsSummary: $("detailsSummary"),
       toast: $("toast"),
     };
 
+    function escapeHtml(value) {
+      return String(value ?? "").replace(/[&<>"']/g, (char) => ({
+        "&": "&amp;",
+        "<": "&lt;",
+        ">": "&gt;",
+        '"': "&quot;",
+        "'": "&#039;",
+      }[char]));
+    }
     function setStatus(text, icon = "circle-dot") {
       els.statusText.innerHTML = `<i data-lucide="${icon}"></i><span>${text}</span>`;
       lucide.createIcons();
@@ -385,12 +428,13 @@
       els.sceneText.classList.add("empty");
       els.rawText.textContent = "";
       els.runStatus.textContent = "Idle";
+      els.detailsSummary.textContent = "Idle";
     }
-    async function selectExample(index) {
+    async function selectExample(index, options = {}) {
+      const { autoAnalyze = true } = options;
       const example = state.examples[index];
       if (!example) return;
       state.exampleIndex = index;
-      state.exampleMatchesPreset = true;
       for (const button of els.exampleList.querySelectorAll(".example-btn")) {
         button.classList.toggle("active", Number(button.dataset.index) === index);
       }
@@ -400,11 +444,13 @@
       els.tokenValue.textContent = String(example.tokens);
       els.tempInput.value = example.temperature;
       els.tempValue.textContent = Number(example.temperature).toFixed(1);
+      state.exampleMatchesPreset = true;
       resetAnalysis();
       const response = await fetch(example.url);
       const blob = await response.blob();
       const file = new File([blob], example.file, { type: blob.type || "video/mp4" });
       setVideoFile(file, URL.createObjectURL(blob), example.title);
+      if (autoAnalyze && example.cached) await analyze();
     }
     async function loadExamples() {
       const response = await fetch("/examples.json");
@@ -416,11 +462,14 @@
         button.className = "example-btn";
         button.dataset.index = index;
         button.type = "button";
-        button.innerHTML = `<strong>${example.title}</strong><span>${example.task}${example.cached ? " - cached" : ""}</span>`;
+        const thumb = example.thumbnail
+          ? `<img class="example-thumb" src="${escapeHtml(example.thumbnail)}" alt="" loading="lazy" />`
+          : `<span class="example-thumb"></span>`;
+        button.innerHTML = `${thumb}<span class="example-copy"><strong>${escapeHtml(example.title)}</strong><span>${escapeHtml(example.task)}${example.cached ? " - cached" : ""}</span></span>`;
         button.addEventListener("click", () => selectExample(index));
         els.exampleList.appendChild(button);
       });
-      if (state.examples.length) selectExample(0);
+      if (state.examples.length) await selectExample(0);
     }
     function renderTimeline(payload) {
       els.timebar.querySelectorAll(".segment").forEach((node) => node.remove());
@@ -452,7 +501,7 @@
         row.type = "button";
         row.className = "event-row";
         row.dataset.index = index;
-        row.innerHTML = `<time>${seconds(event.start)} - ${seconds(event.end)}</time><span>${event.description}</span>`;
+        row.innerHTML = `<time>${seconds(event.start)} - ${seconds(event.end)}</time><span>${escapeHtml(event.description)}</span>`;
         row.addEventListener("click", () => seekTo(event.start));
         els.eventsList.appendChild(row);
       }
@@ -465,6 +514,7 @@
       els.sceneText.classList.toggle("empty", !payload.scene && !payload.span);
       els.rawText.textContent = payload.raw || "";
       els.runStatus.textContent = `${payload.status}${payload.cached ? " | cache hit" : ""}`;
+      els.detailsSummary.textContent = payload.cached ? "Cache hit" : "Ready";
       updateActiveCaption();
     }
     function seekTo(time) {
@@ -487,10 +537,10 @@
       }
       if (active >= 0) {
         const event = events[active];
-        els.activeCaption.innerHTML = `<small>${seconds(event.start)} - ${seconds(event.end)}</small>${event.description}`;
+        els.activeCaption.innerHTML = `<small>${seconds(event.start)} - ${seconds(event.end)}</small>${escapeHtml(event.description)}`;
         els.activeCaption.classList.remove("empty");
       } else if (state.analysis?.scene) {
-        els.activeCaption.innerHTML = `<small>Scene</small>${state.analysis.scene}`;
+        els.activeCaption.innerHTML = `<small>Scene</small>${escapeHtml(state.analysis.scene)}`;
         els.activeCaption.classList.remove("empty");
       } else {
         els.activeCaption.textContent = "Run analysis to sync captions with playback.";

pid/__init__.py

@@ -0,0 +1,14 @@
+# 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.

pid/_ext/__init__.py

@@ -0,0 +1,4 @@
+"""
+Copied from cosmos repository, with necessary modifications.
+Original source: https://github.com/nvidia-cosmos/cosmos-predict2.5/
+"""

pid/_ext/imaginaire/__init__.py

@@ -0,0 +1,14 @@
+# 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.

pid/_ext/imaginaire/checkpointer/__init__.py

@@ -0,0 +1,14 @@
+# 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.

pid/_ext/imaginaire/checkpointer/base.py

@@ -0,0 +1,177 @@
+# 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 os
+from abc import ABC, abstractmethod
+from typing import Optional
+
+import torch
+
+from pid._ext.imaginaire.config import CheckpointConfig, JobConfig
+from pid._ext.imaginaire.model import ImaginaireModel
+from pid._ext.imaginaire.utils import callback
+from pid._ext.imaginaire.utils.easy_io import easy_io
+
+
+class AbstractCheckpointer(ABC):
+    """The checkpointer class. Supports checkpoint saving/loading to both local disk or object store."""
+
+    def __init__(
+        self,
+        config_checkpoint: CheckpointConfig,
+        config_job: JobConfig,
+        callbacks: Optional[callback.CallBackGroup] = None,
+    ):
+        """Constructor of the checkpointer.
+
+        Args:
+            config_checkpoint (CheckpointConfig): The config object for the checkpointer.
+        """
+        self.config_checkpoint = config_checkpoint
+        # Set the callback functions.
+        self.callbacks = callbacks
+        self.save_to_object_store = config_checkpoint.save_to_object_store.enabled
+        self.load_from_object_store = config_checkpoint.load_from_object_store.enabled
+
+        # Set checkpoint directories for local and object store paths
+        self._local_dirname = os.path.join(config_job.path_local, "checkpoints")
+        self._object_store_dirname = os.path.join(config_job.path, "checkpoints")
+
+        self.strict_resume = config_checkpoint.strict_resume
+        self.load_path = config_checkpoint.load_path or None
+        self.load_training_state = config_checkpoint.load_training_state
+        self.only_load_scheduler_state = config_checkpoint.only_load_scheduler_state
+        self.save_thread = None
+        self.verbose = config_checkpoint.verbose
+        self.keys_not_to_resume = config_checkpoint.keys_not_to_resume
+        self.broadcast_via_filesystem = config_checkpoint.broadcast_via_filesystem
+        # Create the object store client interface.
+        if config_checkpoint.load_from_object_store.enabled:
+            self.load_s3_backend_key = "_ckpt_s3_loader"
+            easy_io.set_s3_backend(
+                key="_ckpt_s3_loader",
+                backend_args={
+                    "backend": "s3",
+                    "path_mapping": {
+                        "s3://ckpt/": f"s3://{config_checkpoint.load_from_object_store.bucket}/",
+                    },
+                    "s3_credential_path": config_checkpoint.load_from_object_store.credentials,
+                },
+            )
+        else:
+            self.load_s3_backend_key = None
+
+        if config_checkpoint.save_to_object_store.enabled:
+            self.save_s3_backend_key = "_ckpt_s3_saver"
+            easy_io.set_s3_backend(
+                key="_ckpt_s3_saver",
+                backend_args={
+                    "backend": "s3",
+                    "path_mapping": {
+                        "s3://ckpt/": f"s3://{config_checkpoint.save_to_object_store.bucket}/",
+                    },
+                    "s3_credential_path": config_checkpoint.save_to_object_store.credentials,
+                },
+            )
+        else:
+            self.save_s3_backend_key = None
+
+    @abstractmethod
+    def save(
+        self,
+        model: ImaginaireModel,
+        optimizer: torch.optim.Optimizer,
+        scheduler: torch.optim.lr_scheduler.LRScheduler,
+        grad_scaler: torch.amp.GradScaler,
+        iteration: int,
+    ) -> None:
+        pass
+
+    @abstractmethod
+    def load(
+        self,
+        model: ImaginaireModel,
+        optimizer: Optional[torch.optim.Optimizer] = None,
+        scheduler: Optional[torch.optim.lr_scheduler.LRScheduler] = None,
+        grad_scaler: Optional[torch.amp.GradScaler] = None,
+    ) -> int:
+        pass
+
+    @property
+    def save_bucket(self):
+        """Get the bucket name for saving checkpoints."""
+        return self.config_checkpoint.save_to_object_store.bucket if self.save_to_object_store else None
+
+    @property
+    def load_bucket(self):
+        """Get the bucket name for loading checkpoints."""
+        return self.config_checkpoint.load_from_object_store.bucket if self.load_from_object_store else None
+
+    @property
+    def save_dirname(self):
+        return (
+            f"s3://{self.save_bucket}/{self._object_store_dirname}"
+            if self.save_to_object_store
+            else self._local_dirname
+        )
+
+    @property
+    def load_dirname(self):
+        return (
+            f"s3://{self.load_bucket}/{self._object_store_dirname}"
+            if self.load_from_object_store
+            else self._local_dirname
+        )
+
+    def finalize(self) -> None:
+        """Finalize the checkpointer."""
+        if self.save_thread:
+            self.save_thread.join()
+
+    def _read_latest_checkpoint_file(self) -> str | None:
+        """Get the file name of the latest saved checkpoint. If it doesn't exist, return None.
+
+        Returns:
+            checkpoint_file (str | None): file name of the latest saved checkpoint.
+        """
+        checkpoint_file = None
+        checkpoint_path = os.path.join(self.load_dirname, "latest_checkpoint.txt")
+        if easy_io.exists(f"{checkpoint_path}", backend_key=self.load_s3_backend_key):
+            checkpoint_file = easy_io.load(f"{checkpoint_path}", backend_key=self.load_s3_backend_key).strip()
+
+        return checkpoint_file
+
+    def _write_latest_checkpoint_file(self, checkpoint_file: str) -> None:
+        """Track the file name of the latest saved checkpoint.
+
+        Args:
+            checkpoint_file (str): file name of the latest saved checkpoint.
+        """
+        content = f"{checkpoint_file}\n"
+        checkpoint_path = os.path.join(self.save_dirname, "latest_checkpoint.txt")
+        easy_io.dump(
+            content,
+            checkpoint_path,
+            backend_key=self.save_s3_backend_key,
+        )
+
+    def _check_checkpoint_exists(self, checkpoint_path: str) -> None:
+        """If the file checkpoint_path does not exist, raise an error.
+
+        Args:
+            checkpoint_path (str): full path to the checkpoint.
+        """
+        if not easy_io.exists(f"{checkpoint_path}", backend_key=self.load_s3_backend_key):
+            raise FileNotFoundError(f"File not found (object store): {checkpoint_path}")

pid/_ext/imaginaire/checkpointer/dcp.py

@@ -0,0 +1,1062 @@
+# 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.
+
+
+"""
+Distributed checkpoint (DCP) directory structure and storage backends.
+
+The checkpointer saves model state in a sharded format across multiple processes:
+
+self.save_dirname/
+├── iter_000000005/                    # Checkpoint at iteration 5
+│   ├── model/                         # Model state shards
+│   │   ├── __0_0.distcp              # Shard 0 from rank 0
+│   │   └── __1_0.distcp              # Shard 1 from rank 1
+│   ├── optim/                        # Optimizer state shards
+│   │   ├── __0_0.distcp              # Shard 0 from rank 0
+│   │   └── __1_0.distcp              # Shard 1 from rank 1
+│   ├── scheduler/                    # Learning rate scheduler state
+│   │   ├── __0_0.distcp              # Shard 0 from rank 0
+│   │   └── __1_0.distcp              # Shard 1 from rank 1
+│   └── trainer/                      # Additional training state
+│       ├── __0_0.distcp              # Shard 0 from rank 0
+│       └── __1_0.distcp              # Shard 1 from rank 1
+└── latest_checkpoint.txt             # Points to most recent checkpoint folder, e.g. iter_000000005
+
+Storage path format:
+  self.save_dirname = "{config_job.path_local}/checkpoints"
+
+The sharded format enables efficient distributed saving/loading by:
+1. Parallelizing I/O across processes
+2. Reducing memory usage per process
+3. Supporting both local and cloud storage backends
+"""
+
+import enum
+import functools
+import multiprocessing
+import os
+import queue
+import re
+import time
+import warnings
+from collections import namedtuple
+from multiprocessing import get_context
+from typing import Any, Dict, List, Optional, Set, Tuple, Union, cast
+
+import torch
+import torch.distributed
+import torch.distributed.checkpoint as dcp
+from torch import nn
+from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter
+from torch.distributed.checkpoint._storage_utils import _storage_setup
+from torch.distributed.checkpoint.default_planner import DefaultSavePlanner
+from torch.distributed.checkpoint.logger import _dcp_method_logger
+from torch.distributed.checkpoint.state_dict import (
+    StateDictOptions,
+    get_model_state_dict,
+    get_optimizer_state_dict,
+    set_model_state_dict,
+    set_optimizer_state_dict,
+)
+from torch.distributed.checkpoint.stateful import Stateful
+from torch.distributed.checkpoint.storage import StorageReader
+from torch.distributed.checkpoint.utils import _api_bc_check, _DistWrapper, _profile
+from torch.distributed.tensor import distribute_tensor
+
+from pid._ext.imaginaire.checkpointer.base import AbstractCheckpointer
+from pid._ext.imaginaire.checkpointer.s3_filesystem import S3StorageReader, S3StorageWriter
+from pid._ext.imaginaire.config import CheckpointConfig, JobConfig
+from pid._ext.imaginaire.model import ImaginaireModel
+from pid._ext.imaginaire.utils import callback, distributed, log, misc
+from pid._ext.imaginaire.utils.easy_io import easy_io
+
+try:
+    """
+    We override the default load function to skip loadding _extra_state keys created by transformer-engine.
+    """
+    from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner as _DefaultLoadPlanner
+    from torch.distributed.checkpoint.default_planner import (
+        DTensor,
+        LoadPlan,
+        _create_read_items,
+        _version,
+        flatten_state_dict,
+    )
+    from torch.distributed.checkpoint.metadata import Metadata, TensorStorageMetadata
+
+    def create_default_local_load_plan(state_dict: dict[str, Any], metadata: Metadata, strict: bool = True) -> LoadPlan:
+        requests = []
+        """
+        Create the ``LoadPlan`` used by DefaultLoadPlanner.
+
+        It produces one read item per value in ``state_dict`` using the metadata in ``metadata``.
+
+        The default behavior is to match key exactly between state_dict and metadata.
+        It handles resharding by issuing multiple read requests against storage in order to match
+        load requirements.
+        """
+
+        for fqn, obj in state_dict.items():
+            if fqn.endswith("._extra_state"):  # dirty TE attention package!
+                continue
+            # ignore state_dict keys which do not exist in `state_dict` if strict=False
+            if fqn not in metadata.state_dict_metadata:
+                if strict:
+                    raise RuntimeError(f"Missing key in checkpoint state_dict: {fqn}.")
+                else:
+                    log.warning(f"Local load plan: Missing key in checkpoint state_dict: {fqn}.")
+                    continue
+
+            md = metadata.state_dict_metadata[fqn]
+
+            if (
+                isinstance(md, TensorStorageMetadata)
+                and getattr(obj, "size", None) is not None
+                and md.size != obj.size()
+            ):
+                if not strict:
+                    log.critical(f"Size mismatch between saved {md.size} and current: {obj.size()} for {fqn}")
+                    continue
+                else:
+                    raise ValueError(
+                        f"Size mismatch between saved {md.size} and current: {obj.size()} for {fqn}",
+                    )
+
+            # Since DTensor supports submesh, adding extra check to ensure _create_read_items()
+            # gets called only when the current rank is part of the mesh for the corresponding DTensor.
+            if isinstance(obj, DTensor):
+                if obj.device_mesh.get_coordinate() is not None:
+                    requests += _create_read_items(fqn, md, obj)
+            else:
+                requests += _create_read_items(fqn, md, obj)
+
+        return LoadPlan(requests)
+
+    class DefaultLoadPlanner(_DefaultLoadPlanner):
+        def create_local_plan(self) -> LoadPlan:
+            assert self.metadata is not None
+            if self.flatten_state_dict:
+                # To support checkpoints that are saved before v2.4, we have to
+                # differentiate if the missing keys are due to old checkpoints.
+                # The contracts are:
+                # 1. There are 3 cases when we found a missing key.
+                #    1.1 Actual missing key, but allow_partial_load is False
+                #    1.2 Actual missing key, but allow_partial load is True
+                #    1.3 Old checkpoint, but allow_partial_load is False
+                #    1.4 Old checkpoint, but allow_partial_load is True
+                # 2. If we found a missing key, we first convert the keys back to
+                #    the key format of v2.3
+                # 3. If the previous missing keys are in the v2.3 keys, we assume
+                #    this is a old checkpoint.
+                # 4. Pass the state_dict to `create_default_local_load_plan()`,
+                #    which has the logic to check missing for allow_partial_load.
+                # So for 1.2 and 1.4 cases, we delegate allow_partial_load check to
+                # `create_default_local_load_plan()`. The logic here is to determine
+                # whether the checkpoint belong to 2.3 (or before) or 2.4 (or after).
+                current_keys = set(self.state_dict.keys())
+                load_keys = set(self.metadata.state_dict_metadata.keys())
+                missing_keys = load_keys - current_keys
+                if missing_keys:
+                    _version._derived_version = "2_3"
+                    old_state_dict, old_mappings = flatten_state_dict(self.original_state_dict)
+                    old_keys = set(old_state_dict.keys())
+                    if old_keys & missing_keys:
+                        self.state_dict, self.mappings = old_state_dict, old_mappings
+                    # _derived_version is only used by flatten_state_dict now.
+                    # Set it back to None so that later we can save to a new version.
+                    _version._derived_version = None
+
+            return create_default_local_load_plan(self.state_dict, self.metadata, not self.allow_partial_load)
+
+    log.info("for the back comptiable pytorch! New DefaultLoadPlanner class is created.")
+
+    @_dcp_method_logger(log_exceptions=True)
+    @_api_bc_check
+    def load(
+        state_dict: dict[str, Any],
+        *,
+        checkpoint_id: Union[str, os.PathLike, None] = None,
+        storage_reader: Optional[StorageReader] = None,
+        planner: Optional[DefaultLoadPlanner] = None,
+        process_group: Optional[torch.distributed.ProcessGroup] = None,
+        no_dist: bool = False,
+    ) -> None:
+        """
+        We override the default load function to perform a load plan check for mismatched/missing keys.
+        ==========================Original Doc string=====================================
+        Load a checkpoint into a distributed state dict in SPMD style.
+
+        Each rank must have the same keys in their ``state_dict`` provided to this
+        API. Mismatched keys may result in hangs or errors. If unsure, you can use
+        the ``utils._assert_same_keys`` API to check (but may incur communication
+        costs).
+
+        Each rank will try to read the least amount of data necessary
+        to fullfill the requested `state_dict`. When loading :class:`ShardedTensor`
+        or :class:`DTensor` instances, each rank only reads data for their local shards.
+
+        For each ``Stateful`` object (having both a ``state_dict`` and a ``load_state_dict``),
+        load will first call ``state_dict`` before attempting deserialization, followed by
+        ``load_state_dict`` once the deserialization is complete.
+        For each non-``Stateful`` object, load will deserailize the object, and then replace
+        it in the ``state_dict`` with the deserialized object.
+
+        .. warning::
+            All tensors in ``state_dict`` must be allocated on their
+            destination device *prior to* calling this function.
+
+            All non-tensor data is loaded using `torch.load()` and modified in place
+            on state_dict.
+
+        .. warning::
+            Users must call `load_state_dict` on the root module to ensure load
+            pos-processing and non-tensor data properly propagates.
+
+        .. note:
+            If no process group is initialized, this function will assume the intent
+            is to load a checkpoint into the local process. This can be useful in the
+            case of local inference, and when using regular Tensors (as opposed to DTensor
+            or ShardedTensor)
+
+        .. note:
+            Rank 0 is assumed to be the coordinator rank.
+
+        Args:
+            state_dict (Dict[str, Any]): The state_dict to load the checkpoint into.
+            checkpoint_id (Union[str, os.PathLike, None]):
+                The ID of this checkpoint instance. The meaning of the checkpoint_id
+                depends on the storage. It can be a path to a folder or to a file.
+                It can also be a key if the storage is a key-value store.
+                (Default: ``None``)
+            storage_reader (Optional[StorageReader]):
+                Instance of StorageWriter used to perform reads. If this is not
+                specified, DCP will automatically infer the reader based on the
+                checkpoint_id. If checkpoint_id is also None, an exception will
+                be raised. (Default: ``None``)
+            planner (Optional[LoadPlanner]):
+                Instance of LoadPlanner. If this is not specificed, the default
+                planner will be used. (Default: ``None``)
+            process_group (Optional[ProcessGroup]):
+                ProcessGroup to be used for cross-rank synchronization.
+                (Default: ``None``)
+            no_dist (bool): If ``True``, this function will assume the intent is to load
+                a checkpoint without using cross-rank synchronization. (Default: ``False``)
+        Returns:
+            None.
+
+        Examples
+            >>> # xdoctest: +SKIP
+            >>> my_model = MyModule()
+            >>> optimizer = Adagrad(my_model.parameters())
+            >>> model_state_dict = my_model.state_dict()
+            >>> fs_storage_reader = torch.distributed.checkpoint.FileSystemReader(
+            ...     "/checkpoint/1"
+            ... )
+
+            >>> torch.distributed.checkpoint.load_state_dict(
+            >>>     state_dict=model_state_dict,
+            >>>     storage_reader=fs_storage_reader,
+            >>> )
+
+            >>> # module.load_state_dict() function might have customized steps
+            >>> # to flush the state_dict, must call it to
+            >>> # ensure correct behavior.
+            >>> my_model.load_state_dict(model_state_dict)
+
+        .. note::
+            load_state_dict uses collectives to coordinate reads across ranks.
+            For NCCL-based process groups, internal tensor representations of
+            objects must be moved to the GPU device before communication takes place.
+            In this case, the device used is given by ``torch.cuda.current_device()``
+            and it is the user's responsibility to ensure that this is set so that each
+            rank has an individual GPU, via ``torch.cuda.set_device()``.
+        """
+
+        no_dist = no_dist or (not torch.distributed.is_available()) or (not torch.distributed.is_initialized())
+        if no_dist:
+            warnings.warn(
+                "torch.distributed is disabled, unavailable or uninitialized, assuming the intent is to load in a single process."
+            )
+
+        with _profile():
+            storage_reader = cast(StorageReader, _storage_setup(storage_reader, checkpoint_id, reader=True))
+
+            # All ranks must have the same keys in their `state_dict` provided to
+            # this API.  See documentation for more details.
+            # Here we simply sort the keys to ensure that all ranks load values in
+            # the same order.
+            keys = sorted(state_dict.keys())
+
+            statetful_sd = {}
+            for key in keys:
+                if key not in state_dict:
+                    continue
+                elem = state_dict[key]
+                statetful_sd[key] = elem.state_dict() if isinstance(elem, Stateful) else elem
+
+            _load_state_dict(
+                state_dict=statetful_sd,
+                storage_reader=storage_reader,
+                process_group=process_group,
+                no_dist=no_dist,
+                planner=planner,
+            )
+            for key in keys:
+                if key not in state_dict:
+                    continue
+                elem = state_dict[key]
+                if isinstance(elem, Stateful):
+                    # If the state_dict is a Stateful object,
+                    # DCP does an in-place load in the original state dict.
+                    elem.load_state_dict(statetful_sd[key])
+                else:
+                    # Otherwise, replace the state_dict with the loaded state_dict.
+                    state_dict[key] = statetful_sd[key]
+
+    def _load_state_dict(
+        state_dict: dict[str, Any],
+        storage_reader: StorageReader,
+        process_group: Optional[torch.distributed.ProcessGroup] = None,
+        coordinator_rank: int = 0,
+        no_dist: bool = False,
+        planner: Optional[DefaultLoadPlanner] = None,
+    ) -> None:
+        torch._C._log_api_usage_once("torch.distributed.checkpoint.load_state_dict")
+
+        distW = _DistWrapper(process_group, not no_dist, coordinator_rank)
+        if planner is None:
+            planner = DefaultLoadPlanner()
+
+        ckpt_kwargs = {}
+        if (ckpt_id := getattr(storage_reader, "checkpoint_id", None)) is not None:
+            ckpt_kwargs["checkpoint_id"] = ckpt_id
+            ckpt_kwargs["process_group"] = distW.group
+
+        @_dcp_method_logger(**ckpt_kwargs)
+        def local_step():
+            assert planner is not None
+            metadata = storage_reader.read_metadata()
+            planner.set_up_planner(state_dict, metadata, distW.is_coordinator)
+            storage_reader.set_up_storage_reader(metadata, distW.is_coordinator)
+
+            local_plan = planner.create_local_plan()
+            local_plan = storage_reader.prepare_local_plan(local_plan)
+            return local_plan
+
+        @_dcp_method_logger(**ckpt_kwargs)
+        def global_step(all_local_plans):
+            assert planner is not None
+            all_local_plans = planner.create_global_plan(all_local_plans)
+            all_local_plans = storage_reader.prepare_global_plan(all_local_plans)
+            return all_local_plans
+
+        central_plan: LoadPlan = distW.reduce_scatter("plan", local_step, global_step)
+        if distW.is_coordinator:
+            # Compare central_plan items with storage_reader.storage_data keys
+            dest_fqns = set()
+            storage_fqns = set()
+
+            # Extract FQNs from central_plan items
+            for item in central_plan.items:
+                if hasattr(item, "dest_index") and hasattr(item.dest_index, "fqn"):
+                    dest_fqns.add(item.dest_index.fqn)
+                if hasattr(item, "storage_index") and hasattr(item.storage_index, "fqn"):
+                    storage_fqns.add(item.storage_index.fqn)
+
+            # Get storage data keys
+            storage_data_keys = set()
+            if hasattr(storage_reader, "storage_data") and storage_reader.storage_data is not None:
+                storage_data_keys = set(item[0].fqn for item in storage_reader.storage_data.items())
+            state_dict_keys = set(state_dict.keys())
+            # Compare sets and log differences
+            # Remove any item that has "_extra_state" as substring in the sets
+            state_dict_keys = {fqn for fqn in state_dict_keys if "_extra_state" not in fqn}
+            dest_fqns = {fqn for fqn in dest_fqns if "_extra_state" not in fqn}
+            storage_fqns = {fqn for fqn in storage_fqns if "_extra_state" not in fqn}
+            storage_data_keys = {fqn for fqn in storage_data_keys if "_extra_state" not in fqn}
+
+            log.info("=== Load Plan FQN Analysis ===")
+            log.info(f"State Dict FQNs count: {len(state_dict_keys)}")
+            log.info(f"Destination FQNs count (without _extra_state): {len(dest_fqns)}")
+            log.info(f"Loaded FQNs count (without _extra_state): {len(storage_fqns)}")
+            log.info(f"In Storage keys count (without _extra_state): {len(storage_data_keys)}")
+
+            # Find missing keys in each direction
+            state_dict_missing_from_dest = state_dict_keys - dest_fqns
+            storage_data_missing_from_storage_fqns = storage_data_keys - storage_fqns
+
+            if state_dict_missing_from_dest:
+                log.info(
+                    f"State Dict FQNs missing from load plan ({len(state_dict_missing_from_dest)} items): {sorted(state_dict_missing_from_dest)}"
+                )
+            else:
+                log.info("✓ All State Dict FQNs found in storage_data")
+
+            if storage_data_missing_from_storage_fqns:
+                # If there are more than 100 "net_ema" keys in storage_data_missing_from_storage_fqns, summarize them
+                net_ema_keys = {k for k in storage_data_missing_from_storage_fqns if "net_ema" in k}
+                if len(net_ema_keys) > 100:
+                    storage_data_missing_from_storage_fqns = storage_data_missing_from_storage_fqns - net_ema_keys
+                    storage_data_missing_from_storage_fqns = set(
+                        storage_data_missing_from_storage_fqns
+                    )  # ensure set type
+                    storage_data_missing_from_storage_fqns.add("net_ema")
+                    log.info(
+                        f"Summarized {len(net_ema_keys)} 'net_ema' keys as 'net_ema' in missing storage data keys."
+                    )
+                log.info(
+                    f"Storage data keys not loaded by load plan ({len(storage_data_missing_from_storage_fqns)} items): {sorted(storage_data_missing_from_storage_fqns)}"
+                )
+            else:
+                log.info("✓ All storage data keys found in Loaded FQNs")
+
+            log.info("=== End Load Plan FQN Analysis ===")
+
+        @_dcp_method_logger(**ckpt_kwargs)
+        def read_data():
+            assert planner is not None
+            final_local_plan = planner.finish_plan(central_plan)
+            all_reads = storage_reader.read_data(final_local_plan, planner)
+
+            all_reads.wait()
+            return None
+
+        _ = distW.all_gather("read", read_data)
+
+except ImportError as e:
+    from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner
+
+    log.critical(f"{e}, using default planner")
+
+StateDictItemPath = namedtuple("StateDictItemPath", ["state_dict", "save_path"])
+
+
+class ModelWrapper(Stateful):
+    """Wrapper for model state dict handling with DCP (Distributed Checkpoint).
+
+    Usage flow: state_dict() -> dcp.load() -> load_state_dict()
+
+    When load_ema_to_reg=True, the three steps work together to load EMA weights from
+    checkpoint into the regular model:
+      1. state_dict():      prepare keys so dcp.load() reads EMA weights from checkpoint
+      2. dcp.load():        fills the state_dict with checkpoint data (in-place)
+      3. load_state_dict(): remap loaded EMA weights back to regular model keys, then apply
+    """
+
+    def __init__(self, model: Union[nn.Module, List[nn.Module]], load_ema_to_reg: bool = False):
+        self.model = [model] if isinstance(model, nn.Module) else model
+        self.load_ema_to_reg = load_ema_to_reg
+        if self.load_ema_to_reg:
+            assert isinstance(model, ImaginaireModel), (
+                f"replicate_ema_to_reg requires an ImaginaireModel, got {type(model).__name__}"
+            )
+
+    def state_dict(self) -> Dict[str, Any]:
+        """Build a state_dict whose keys will be matched against checkpoint keys by dcp.load().
+
+        When load_ema_to_reg=True, we manipulate the keys so that dcp.load() reads the
+        checkpoint's net_ema.* weights (instead of net.*):
+
+        Case A - ema.enabled=False (model has no EMA module, only net.*):
+          Rename net.* -> net_ema.* so dcp.load() matches checkpoint's net_ema.* keys.
+          Result after dcp.load(): state_dict has net_ema.* filled with checkpoint EMA weights.
+
+        Case B - ema.enabled=True (model has both net.* and net_ema.*):
+          Drop net.* keys entirely, keep net_ema.* (and other keys like fake_score.*).
+          Result after dcp.load(): net_ema.* filled from checkpoint; other keys (e.g. fake_score.*)
+          loaded if present in checkpoint, otherwise unchanged.
+        """
+        _state_dict = {k: v for sd in map(get_model_state_dict, self.model) for k, v in sd.items()}
+        if self.load_ema_to_reg:
+            if not self.model[0].config.ema.enabled:
+                # Case A: model only has net.* -> rename to net_ema.* to match checkpoint EMA keys.
+                # Auxiliary keys (aux_losses.*, teacher_align_loss.*, etc.) are passed through as-is.
+                all_keys = list(_state_dict.keys())
+                for k in all_keys:
+                    if k.startswith("net."):
+                        _state_dict[k.replace("net.", "net_ema.", 1)] = _state_dict.pop(k)
+            else:
+                # Case B: model has net.* + net_ema.* -> drop net.*, only load net_ema.* from ckpt
+                log.warning("EMA is enabled, will only load EMA weights from checkpoint file.")
+                all_keys = list(_state_dict.keys())
+                for k in all_keys:
+                    if k.startswith("net_ema."):
+                        break
+                else:
+                    raise ValueError("No EMA keys found in state_dict")
+                _state_dict = {k: _state_dict[k] for k in all_keys if not k.startswith("net.")}
+
+        if hasattr(self.model[0].config, "lora_config") and self.model[0].config.lora_config.enabled:
+            """
+            When using LoRA, `inject_adapter_in_model` modifies the target modules in place.
+            For example, `blocks[0].attn.q_proj.weight` will be modified to `blocks[0].attn.q_proj.base_layer.weight`.
+            This means that the model will have the key `blocks[0].attn.q_proj.base_layer.weight`,
+            but the checkpoint will have the key `blocks[0].attn.q_proj.weight`.
+            We need to map the model key to the checkpoint key.
+            """
+            self.checkpoint_to_model_key = {}
+            mapping_keys = {"base_layer.": "", "base_model.model.": ""}
+            keys_to_update = []
+            for k in _state_dict.keys():
+                new_key = k
+                for from_key, to_key in mapping_keys.items():
+                    new_key = new_key.replace(from_key, to_key)
+                if new_key != k:
+                    keys_to_update.append((k, new_key))
+                    self.checkpoint_to_model_key[new_key] = k
+            for k, new_key in keys_to_update:
+                _state_dict[new_key] = _state_dict.pop(k)
+
+        return _state_dict
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        """Apply the checkpoint-loaded state_dict back to the model.
+
+        When load_ema_to_reg=True, remap the loaded EMA weights to regular model keys:
+
+        Case A - ema.enabled=False:
+          state_dict has net_ema.* (loaded from checkpoint) -> rename back to net.* for the model.
+
+        Case B - ema.enabled=True:
+          state_dict has net_ema.* (loaded from checkpoint) -> clone each net_ema.* as net.*
+          so the regular model gets the same EMA weights. net_ema.* is also kept.
+        """
+        if hasattr(self.model[0].config, "lora_config") and self.model[0].config.lora_config.enabled:
+            if hasattr(self, "checkpoint_to_model_key"):
+                for checkpoint_key, model_key in self.checkpoint_to_model_key.items():
+                    state_dict[model_key] = state_dict.pop(checkpoint_key)
+            else:
+                raise ValueError("checkpoint_to_model_key is not set by `state_dict`")
+
+        if self.load_ema_to_reg:
+            if not self.model[0].config.ema.enabled:
+                # Case A: rename net_ema.* back to net.* so set_model_state_dict can apply them.
+                # Auxiliary keys (aux_losses.*, teacher_align_loss.*, etc.) are passed through as-is.
+                all_keys = list(state_dict.keys())
+                for k in all_keys:
+                    if k.startswith("net_ema."):
+                        state_dict[k.replace("net_ema.", "net.", 1)] = state_dict.pop(k)
+            else:
+                # Case B: clone net_ema.* -> net.* so regular model gets EMA weights too.
+                # Auxiliary keys are kept unchanged.
+                log.warning("EMA is enabled, will load EMA weights to regular model weights")
+                all_keys = list(state_dict.keys())
+                for k in all_keys:
+                    if k.startswith("net_ema."):
+                        state_dict[k.replace("net_ema.", "net.", 1)] = torch.clone(state_dict[k])
+        func = functools.partial(
+            set_model_state_dict,
+            model_state_dict=state_dict,
+            options=StateDictOptions(strict=False),
+        )
+        list(map(func, self.model))
+
+
+class OptimizerWrapper(Stateful):
+    """
+    Wrapper for optimizer state dict handling with DCP (Distributed Checkpoint).
+
+    Note: For multi-optimizer models (e.g., distillation), use DistillationCheckpointer
+    from pid._src.checkpointer.dcp_distill which has fallback handling for
+    FSDP2 parameter ID mismatches.
+    """
+
+    def __init__(
+        self,
+        model: Union[nn.Module, List[nn.Module]],
+        optim: Union[torch.optim.Optimizer, List[torch.optim.Optimizer]],
+    ) -> None:
+        self.model = [model] if isinstance(model, nn.Module) else model
+        self.optim = [optim] if isinstance(optim, torch.optim.Optimizer) else optim
+
+    def state_dict(self) -> Dict[str, Any]:
+        func = functools.partial(
+            get_optimizer_state_dict,
+            options=StateDictOptions(flatten_optimizer_state_dict=True),
+        )
+        return {k: v for sd in map(func, self.model, self.optim) for k, v in sd.items()}
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        func = functools.partial(
+            set_optimizer_state_dict,
+            optim_state_dict=state_dict,
+            options=StateDictOptions(flatten_optimizer_state_dict=True),
+        )
+        list(map(func, self.model, self.optim))
+
+
+class AsyncMode(str, enum.Enum):
+    DISABLED = "disabled"
+    ASYNC_WITH_PINNED_MEM = "async_with_pinned_mem"
+
+
+class Terminate:
+    pass
+
+
+class SaveDone:
+    def __init__(self, iteration: int, elapsed_time: float, succeeded: bool):
+        self.iteration = iteration
+        self.elapsed_time = elapsed_time
+        self.succeeded = succeeded
+
+    def __str__(self):
+        return f"SaveDone(iteration={self.iteration}, elapsed_time={self.elapsed_time}, succeeded={self.succeeded})"
+
+
+def save_checkpoint_in_background(
+    receiver_queue: multiprocessing.Queue,
+    sender_queue: multiprocessing.Queue,
+    checkpoint_config: CheckpointConfig,
+    job_config: JobConfig,
+) -> None:
+    """
+    Handles model checkpoint saving in a separate background process using PyTorch's distributed functionality.
+    This function runs in a dedicated process to avoid blocking the main training loop.
+
+    Args:
+        receiver_queue: Queue to receive state dictionaries and commands from the main process
+        sender_queue: Queue to send completion signals back to the main process
+        checkpoint_config: Configuration settings for checkpoint saving behavior
+        job_config: Configuration settings for the training job
+
+    Flow:
+        1. Initializes distributed processing environment
+        2. Continuously waits for state dictionaries to save
+        3. Saves checkpoints asynchronously
+        4. Signals completion back to main process
+        5. Terminates when receiving a Terminate signal
+
+    Raises:
+        AssertionError: If received object is neither Terminate signal nor valid state dict tuple
+
+    Note:
+        - Uses a different port than the main process to avoid conflicts
+        - Disables TorchElastic agent store for checkpoint operations
+        - Automatically cleans up distributed process group on exit
+    """
+    # Configure distributed environment
+    os.environ["MASTER_PORT"] = str(int(os.environ["MASTER_PORT"]) + 2)
+    os.environ["TORCHELASTIC_USE_AGENT_STORE"] = "False"
+
+    # Set up GPU device and distributed processing
+    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+    distributed.init()
+
+    # Initialize checkpointing mechanism
+    checkpoint_handler = DistributedCheckpointer(checkpoint_config, job_config, None, disable_async=True)
+
+    try:
+        while True:
+            log.debug("Checkpoint background process is ready for next task, waiting for new state_dict")
+            received_data = receiver_queue.get()
+            log.debug("Received new state_dict")
+
+            if isinstance(received_data, Terminate):
+                log.info("Received termination signal in checkpoint background process, closing sender queue")
+                sender_queue.put(Terminate())
+                sender_queue.close()
+                return
+
+            assert isinstance(received_data, tuple), "Received data must be a tuple of (state_dict, checkpoint_path)"
+            state_dict, checkpoint_path = received_data
+
+            # Save checkpoint and measure time taken
+            start_time = time.monotonic()
+            iteration = state_dict["trainer"][0]["iteration"]
+            elapsed_time = 0
+            succeeded = False
+            try:
+                checkpoint_handler.save_state_dict_worker(state_dict, checkpoint_path)
+                elapsed_time = time.monotonic() - start_time
+                log.info(
+                    f"Checkpoint saved successfully in background process. Time taken: {elapsed_time:.2f} seconds, iteration: {iteration}"
+                )
+                succeeded = True
+            except Exception as e:
+                log.error(f"Error saving checkpoint to {checkpoint_path}: {e}")
+                # continue because if the thread exits, the main thread keeps on adding to the queue
+            finally:
+                if elapsed_time == 0:
+                    elapsed_time = time.monotonic() - start_time
+                sender_queue.put(SaveDone(iteration, elapsed_time, succeeded))
+
+    finally:
+        log.info("Cleaning up: destroying distributed process group")
+        torch.distributed.destroy_process_group()
+
+
+class DistributedCheckpointer(AbstractCheckpointer):
+    KEYS_TO_SAVE = ["model", "optim", "scheduler", "trainer"]
+
+    def __init__(
+        self,
+        config_checkpoint: CheckpointConfig,
+        config_job: JobConfig,
+        callbacks: Optional[callback.CallBackGroup] = None,
+        disable_async: bool = False,
+    ):
+        super().__init__(config_checkpoint, config_job, callbacks)
+        self.config_checkpoint = config_checkpoint
+        if config_checkpoint.dcp_async_mode_enabled:
+            self.async_mode = AsyncMode.ASYNC_WITH_PINNED_MEM
+        else:
+            self.async_mode = AsyncMode.DISABLED
+
+        if disable_async:
+            self.async_mode = AsyncMode.DISABLED
+
+        if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
+            ctx = get_context("spawn")
+            self.mp_queue_send = ctx.Queue()
+            self.mp_queue_recv = ctx.Queue()
+            self.mp = ctx.Process(
+                target=save_checkpoint_in_background,
+                args=(
+                    self.mp_queue_send,
+                    self.mp_queue_recv,
+                    config_checkpoint,
+                    config_job,
+                ),
+                daemon=True,
+            )
+            self.mp.start()
+            self.cpu_offload_state_dict = None
+            self.staging = False
+            self.staging_ckpt_file = None
+            self.staging_stream = torch.cuda.Stream()
+
+    def keys_to_resume_during_load(self) -> Tuple[Set, Union[str, None]]:
+        latest_checkpoint_file = self._read_latest_checkpoint_file()
+
+        resume_keys = []
+
+        if latest_checkpoint_file is not None:
+            # 1. Resume training from latest_checkpoint.txt under the same name.
+            checkpoint_path = os.path.join(self.load_dirname, latest_checkpoint_file)
+            resume_keys.extend(self.KEYS_TO_SAVE)
+        else:
+            if self.load_path and not self.load_path.endswith(".pth"):
+                # 2. Load the module weights specified by config_checkpoint.path.
+                checkpoint_path = self.load_path
+                if self.load_s3_backend_key:
+                    checkpoint_path = f"s3://{self.config_checkpoint.load_from_object_store.bucket}/{checkpoint_path}"
+                    if not re.search(r"/checkpoints/iter_\d{9}/?$", checkpoint_path):
+                        old_ckpt_path = checkpoint_path
+                        # If path doesn't end with specific checkpoint, read latest checkpoint file
+                        latest_ckpt_path = os.path.join(checkpoint_path, "checkpoints/latest_checkpoint.txt")
+                        if easy_io.exists(latest_ckpt_path, backend_key=self.load_s3_backend_key):
+                            checkpoint_file = easy_io.load(
+                                latest_ckpt_path, backend_key=self.load_s3_backend_key
+                            ).strip()
+                            checkpoint_path = f"{checkpoint_path}/checkpoints/{checkpoint_file}"
+                        else:
+                            log.warning(
+                                f"Latest checkpoint file {latest_ckpt_path} not found, load from {old_ckpt_path}"
+                            )
+                            checkpoint_path = old_ckpt_path
+                if self.load_training_state:
+                    resume_keys.extend(self.KEYS_TO_SAVE)
+                else:
+                    resume_keys.append("model")
+                    if self.only_load_scheduler_state:
+                        resume_keys.append("scheduler")
+            elif self.load_path and self.load_path.endswith(".pth"):
+                checkpoint_path = self.load_path
+            else:
+                checkpoint_path = None
+        if len(self.keys_not_to_resume) > 0:
+            for key in self.keys_not_to_resume:
+                assert key in self.KEYS_TO_SAVE, f"Invalid key to resume: {key} not in {self.KEYS_TO_SAVE}"
+            resume_keys = [key for key in resume_keys if key not in self.keys_not_to_resume]
+        is_self_resume = latest_checkpoint_file is not None
+        return set(resume_keys), checkpoint_path, is_self_resume
+
+    @misc.timer("checkpoint loading")
+    def load(
+        self,
+        model: ImaginaireModel,
+        optimizer: torch.optim.Optimizer | None = None,
+        scheduler: torch.optim.lr_scheduler.LRScheduler | None = None,
+        grad_scaler: torch.amp.GradScaler | None = None,
+    ) -> int:
+        if self.callbacks is not None:
+            self.callbacks.on_load_checkpoint_start(model)
+
+        resume_keys, checkpoint_path, is_self_resume = self.keys_to_resume_during_load()
+        resume_keys = sorted(resume_keys)
+        log.info(f"Resuming ckpt {checkpoint_path} with keys: {resume_keys}")
+
+        iteration = 0
+
+        if checkpoint_path is not None and not checkpoint_path.endswith(".pth"):
+            self._check_checkpoint_exists(checkpoint_path)
+            for key in resume_keys:
+                load_planner = DefaultLoadPlanner(allow_partial_load=True)
+                cur_key_ckpt_full_path = os.path.join(checkpoint_path, key)
+                log.info(f"Start loading checkpoint from {checkpoint_path}")
+                storage_reader = self.get_storage_reader(cur_key_ckpt_full_path)
+                torch.distributed.barrier()
+                log.info(f"starting {cur_key_ckpt_full_path}", rank0_only=False)
+                if key == "model":
+                    log.info("- Loading the model...")
+                    should_replicate_ema = (
+                        self.config_checkpoint.replicate_ema_to_reg_in_training
+                        and not is_self_resume
+                        and not self.load_training_state
+                    )
+                    if should_replicate_ema:
+                        ckpt_metadata = storage_reader.read_metadata()
+                        ckpt_keys = set(ckpt_metadata.state_dict_metadata.keys())
+                        if not any(k.startswith("net_ema.") for k in ckpt_keys):
+                            raise ValueError(
+                                f"replicate_ema_to_reg_in_training is True, but checkpoint at "
+                                f"{cur_key_ckpt_full_path} does not contain EMA weights (no 'net_ema.*' keys). "
+                                f"Set replicate_ema_to_reg_in_training=False or use a checkpoint with EMA weights."
+                            )
+                    _model_wrapper = ModelWrapper(model, load_ema_to_reg=should_replicate_ema)
+                    _state_dict = _model_wrapper.state_dict()
+                    load(_state_dict, storage_reader=storage_reader, planner=load_planner)
+                    _model_wrapper.load_state_dict(_state_dict)
+                elif key == "optim":
+                    log.info("- Loading the optimizer...")
+                    # Note: For multi-optimizer models (e.g., distillation), use DistillationCheckpointer
+                    # from pid._src.checkpointer.dcp_distill instead.
+                    _optim_wrapper = OptimizerWrapper(model, optimizer)
+                    _state_dict = _optim_wrapper.state_dict()
+                    dcp.load(
+                        _state_dict,
+                        storage_reader=storage_reader,
+                        planner=load_planner,
+                    )
+                    _optim_wrapper.load_state_dict(_state_dict)
+                elif key == "scheduler":
+                    log.info("- Loading the scheduler...")
+                    # Note: For multi-scheduler models (e.g., distillation), use DistillationCheckpointer
+                    # from pid._src.checkpointer.dcp_distill instead.
+                    _state_dict = scheduler.state_dict()
+                    dcp.load(
+                        _state_dict,
+                        storage_reader=storage_reader,
+                        planner=load_planner,
+                    )
+                    scheduler.load_state_dict(_state_dict)
+                elif key == "trainer":
+                    log.info("- Loading the trainer...")
+                    _state_dict = {
+                        "grad_scaler": grad_scaler.state_dict(),
+                        "iteration": iteration,
+                    }
+                    dcp.load(
+                        _state_dict,
+                        storage_reader=storage_reader,
+                        planner=load_planner,
+                    )
+                    grad_scaler.load_state_dict(_state_dict["grad_scaler"])
+                    iteration = _state_dict["iteration"]
+                else:
+                    raise ValueError(f"Invalid key: {key}. not support to resume.")
+            if self.callbacks is not None:
+                self.callbacks.on_load_checkpoint(model, state_dict=_state_dict)
+            log.info(f"Loaded checkpoint from {checkpoint_path} in iteration {iteration}")
+        elif checkpoint_path is not None and checkpoint_path.endswith(".pth"):
+            state = easy_io.load(checkpoint_path)
+            model_state = model.net.state_dict()
+
+            for k, v in list(state.items()):
+                tgt = model_state.get(k, None)
+                if tgt is None:
+                    continue
+                # If target param/buffer is a DTensor and checkpoint value is not, distribute it
+                if isinstance(tgt, DTensor) and not isinstance(v, DTensor):
+                    # Match device, dtype, and placements from the target DTensor
+                    v = v.to(tgt.device, dtype=tgt.dtype, copy=False)
+                    v = distribute_tensor(v, tgt.device_mesh, tgt.placements)
+                    state[k] = v
+                # If target is a plain Tensor but checkpoint is a DTensor, bring it local
+                if not isinstance(tgt, DTensor) and isinstance(v, DTensor):
+                    state[k] = v.to_local().to(tgt.device, dtype=tgt.dtype, copy=False)
+
+            model.load_state_dict(state, strict=False, pretrain_copy=True)
+            # Clear unused reserved memory from fp32
+            torch.cuda.empty_cache()
+            log.critical(f"Loaded checkpoint from {checkpoint_path}. **** This only happen at iteration 0 **** ")
+        else:
+            log.info("Training from scratch.")
+        torch.cuda.empty_cache()
+
+        if self.callbacks is not None:
+            self.callbacks.on_load_checkpoint_end(model, iteration=iteration, checkpoint_path=checkpoint_path)
+        return iteration
+
+    def _async_with_pinned_memory(self, checkpoint_file: str, state_dict: Dict[str, Tuple[Any, str]]) -> None:
+        try:
+            from torch.distributed._state_dict_utils import _copy_state_dict, _create_cpu_state_dict
+        except ImportError as e:
+            raise ImportError(
+                "Please install the latest PyTorch nightly to use async checkpointing with pinned memory."
+            ) from e
+        if self.cpu_offload_state_dict is None:
+            log.debug(f"Preparing the CPU memory, {time.monotonic()=}.:.2f")
+            self.cpu_offload_state_dict = _create_cpu_state_dict(state_dict, pin_memory=True, share_memory=True)
+
+        log.debug(f"Staging the state_dict, {time.monotonic()=}.:.2f")
+        with torch.cuda.stream(self.staging_stream):
+            self.cpu_offload_state_dict = _copy_state_dict(
+                state_dict,
+                self.cpu_offload_state_dict,
+                non_blocking=True,
+            )
+            self.staging = True
+            self.staging_ckpt_file = checkpoint_file
+
+        self.maybe_wait_for_staging()
+
+    def maybe_wait_for_staging(self) -> None:
+        if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM and self.staging:
+            if not self.staging_stream.query():
+                self.staging_stream.synchronize()
+
+            def sync_func():
+                self.mp_queue_send.put_nowait((self.cpu_offload_state_dict, self.staging_ckpt_file))
+
+            sync_func()
+            self.staging = False
+
+    def get_previous_checkpoint_results(self, wait_for: int = 0) -> None:
+        """Get the results of previously submitted checkpoints and pass them to callbacks if checkpoint succeeded"""
+        if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
+            try:
+                start_time = time.monotonic()
+                while not self.mp_queue_recv.empty() or wait_for > 0:
+                    try:
+                        ret = self.mp_queue_recv.get(timeout=1)
+                        if isinstance(ret, Terminate):
+                            log.info("Received termination event from checkpoint background process")
+                            break
+                        save_done: SaveDone = ret
+                        log.logger.info(f"Received checkpoint save result: {save_done}")
+                        if self.callbacks is not None and save_done.succeeded:
+                            self.callbacks.on_save_checkpoint_success(
+                                iteration=save_done.iteration, elapsed_time=save_done.elapsed_time
+                            )
+                    except queue.Empty:
+                        elapsed_time = time.monotonic() - start_time
+                        if elapsed_time > wait_for:
+                            break
+            except (EOFError, BrokenPipeError):
+                log.info("Queue was closed by checkpoint background process")
+
+    def get_storage_writer(self, checkpoint_path: str) -> Union[S3StorageWriter, FileSystemWriter]:
+        if self.save_to_object_store:
+            return S3StorageWriter(
+                credential_path=self.config_checkpoint.save_to_object_store.credentials,
+                path=checkpoint_path,
+            )
+        return FileSystemWriter(path=checkpoint_path)
+
+    def get_storage_reader(self, checkpoint_path: str) -> Union[S3StorageReader, FileSystemReader]:
+        if self.load_from_object_store:
+            return S3StorageReader(
+                credential_path=self.config_checkpoint.load_from_object_store.credentials,
+                path=checkpoint_path,
+            )
+        return FileSystemReader(checkpoint_path)
+
+    def save_state_dict_worker(self, to_save_dict: Dict[str, Tuple[Any, str]], checkpoint_file: str) -> None:
+        for k, (v, full_checkpoint_path) in to_save_dict.items():
+            storage_writer = self.get_storage_writer(full_checkpoint_path)
+            dcp.save(
+                v,
+                storage_writer=storage_writer,
+                planner=DefaultSavePlanner(dedup_save_to_lowest_rank=True),
+            )
+
+        if distributed.is_rank0():
+            print(f"Saving last checkpoint file {checkpoint_file}")
+            self._write_latest_checkpoint_file(checkpoint_file)
+
+        log.critical(f"Saved checkpoint to {os.path.join(self.save_dirname, checkpoint_file)}", rank0_only=True)
+
+    def save(
+        self,
+        model: ImaginaireModel,
+        optimizer: torch.optim.Optimizer,
+        scheduler: torch.optim.lr_scheduler.LRScheduler,
+        grad_scaler: torch.amp.GradScaler,
+        iteration: int,
+    ) -> None:
+        """Save network weights, optimizer parameters, scheduler parameters to a checkpoint.
+
+        Note: For multi-optimizer models (e.g., distillation), use DistillationCheckpointer
+        from pid._src.checkpointer.dcp_distill instead.
+
+        Args:
+            model (ImaginaireModel): The PyTorch model.
+            optimizer (torch.optim.Optimizer): The model optimizer.
+            scheduler (torch.optim.lr_scheduler.LRScheduler): The optimization scheduler.
+            grad_scaler (torch.amp.GradScaler): The gradient scaler (for mixed precision training).
+            iteration (int): Current iteration number.
+        """
+        if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
+            self.get_previous_checkpoint_results(wait_for=0)
+
+        if self.callbacks is not None:
+            self.callbacks.on_save_checkpoint_start(model, iteration)
+
+        checkpoint_file = f"iter_{iteration:09}"
+
+        to_save_dict = {
+            "model": ModelWrapper(model).state_dict(),
+            "optim": OptimizerWrapper(model, optimizer).state_dict(),
+            "scheduler": scheduler.state_dict(),
+            "trainer": {
+                "grad_scaler": grad_scaler.state_dict(),
+                "iteration": iteration,
+            },
+        }
+
+        for key in to_save_dict.keys():
+            output_dirname = os.path.join(self.save_dirname, f"iter_{iteration:09}/{key}")
+            to_save_dict[key] = (to_save_dict[key], output_dirname)
+
+        if self.callbacks is not None:
+            self.callbacks.on_save_checkpoint(model, state_dict=to_save_dict)
+
+        if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
+            self._async_with_pinned_memory(checkpoint_file, to_save_dict)
+        else:
+            start_time = time.monotonic()
+            try:
+                self.save_state_dict_worker(to_save_dict, checkpoint_file)
+            finally:
+                if self.callbacks is not None:
+                    self.callbacks.on_save_checkpoint_success(
+                        iteration=iteration, elapsed_time=time.monotonic() - start_time
+                    )
+
+        # This measures exposed (synchronous) checkpoint time, on_save_checkpoint_success()
+        # is instead called to measure the entire duration for asynchronous checkpoint for the async case too.
+        if self.callbacks is not None:
+            self.callbacks.on_save_checkpoint_end(model=None, iteration=iteration)
+
+    def finalize(self) -> None:
+        super().finalize()
+        if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
+            if self.mp and self.mp.is_alive():
+                self.mp_queue_send.put(Terminate())
+                self.get_previous_checkpoint_results(wait_for=60)
+                self.mp.join()

pid/_ext/imaginaire/checkpointer/dummy.py

@@ -0,0 +1,47 @@
+# 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.
+
+from typing import Optional
+
+import torch
+import torch.distributed
+
+from pid._ext.imaginaire.checkpointer.base import AbstractCheckpointer
+from pid._ext.imaginaire.model import ImaginaireModel
+
+
+class Checkpointer(AbstractCheckpointer):
+    """
+    A dummy checkpointer that does not save or load anything. This is useful for debugging jobs or share workload with collobrators.
+    """
+
+    def save(
+        self,
+        model: ImaginaireModel,
+        optimizer: torch.optim.Optimizer,
+        scheduler: torch.optim.lr_scheduler.LRScheduler,
+        grad_scaler: torch.amp.GradScaler,
+        iteration: int,
+    ) -> None:
+        pass
+
+    def load(
+        self,
+        model: ImaginaireModel,
+        optimizer: Optional[torch.optim.Optimizer] = None,
+        scheduler: Optional[torch.optim.lr_scheduler.LRScheduler] = None,
+        grad_scaler: Optional[torch.amp.GradScaler] = None,
+    ) -> int:
+        return 0

pid/_ext/imaginaire/checkpointer/s3_filesystem.py

@@ -0,0 +1,323 @@
+# 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 io
+import json
+import os
+import time
+from contextlib import contextmanager
+from typing import Generator, Union
+from urllib.parse import urlparse
+
+import boto3
+from botocore.config import Config
+from botocore.exceptions import ClientError
+from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter
+from torch.distributed.checkpoint.filesystem import FileSystemBase
+
+from pid._ext.imaginaire.utils import log
+
+
+class S3Stream(io.BytesIO):
+    """
+    Workaround for PyTorch manually closing the stream before we can upload it to S3. We override the close() as noop
+    and instead call our own _true_close() method to close the stream after we are done using it.
+    The commit at fault is https://github.com/pytorch/pytorch/commit/9c909bf3bb122db2cce95e2eb7459bbe50dfa15a
+    """
+
+    def close(self):
+        self.flush()
+        # No close
+
+    def _true_close(self):
+        super().close()
+
+
+class S3FileSystem(FileSystemBase):
+    """Implementation of FileSystemBase for AWS S3 storage."""
+
+    def __init__(
+        self,
+        credential_path: str,
+        max_attempts: int = 20,
+        initial_backoff: float = 1.0,
+        max_backoff: float = 30.0,
+        backoff_factor: float = 2.0,
+    ) -> None:
+        """
+        Initialize S3FileSystem with retry configuration.
+
+        Args:
+            credential_path: Path to AWS credentials JSON file
+            max_attempts: Maximum number of retry attempts
+            initial_backoff: Initial backoff time in seconds
+            max_backoff: Maximum backoff time in seconds
+            backoff_factor: Multiplicative factor for backoff time
+        """
+        with open(credential_path, "r") as f:
+            conf = json.load(f)
+
+        # Configure boto3 with retry settings
+        config = Config(
+            retries=dict(max_attempts=max_attempts, mode="adaptive"),  # Adaptive mode automatically handles throttling
+            connect_timeout=60,
+            read_timeout=60,
+            request_checksum_calculation="when_required",  # Data integrity check for uploads and downloads
+            response_checksum_validation="when_required",  # Data integrity check for uploads and downloads
+        )
+
+        self.s3_client = boto3.client("s3", config=config, **conf)
+        self.max_attempts = max_attempts
+        self.initial_backoff = initial_backoff
+        self.max_backoff = max_backoff
+        self.backoff_factor = backoff_factor
+
+    def _retry_with_backoff(self, operation_func, *args, **kwargs):
+        """
+        Execute an operation with exponential backoff retry logic.
+
+        Args:
+            operation_func: Function to execute
+            *args: Positional arguments for the function
+            **kwargs: Keyword arguments for the function
+
+        Returns:
+            Result of the operation function
+
+        Raises:
+            Exception: If all retry attempts fail
+        """
+        last_exception = None
+        backoff = self.initial_backoff
+
+        for attempt in range(self.max_attempts):
+            try:
+                return operation_func(*args, **kwargs)
+            except ClientError as e:
+                error_code = e.response.get("Error", {}).get("Code", "")
+                log.info(f"S3 Filesystem: Received ClientError: {error_code}", rank0_only=False)
+
+                # Handle specific error cases
+                if error_code in ["SlowDown", "ThrottlingException", "RequestLimitExceeded", "InternalError"]:
+                    last_exception = e
+                    if attempt < self.max_attempts - 1:  # Don't sleep on last attempt
+                        current_backoff = min(backoff, self.max_backoff)
+                        log.info(f"S3 Filesystem: Retrying in {current_backoff} seconds", rank0_only=False)
+                        time.sleep(current_backoff)
+                        backoff *= self.backoff_factor
+                        continue
+                # For other client errors, raise immediately
+                raise
+            except Exception as e:
+                log.info(f"S3 Filesystem: Received Exception: {str(e)}", rank0_only=False)
+                last_exception = e
+                if attempt < self.max_attempts - 1:
+                    current_backoff = min(backoff, self.max_backoff)
+                    log.info(f"S3 Filesystem: Retrying in {current_backoff} seconds", rank0_only=False)
+                    time.sleep(current_backoff)
+                    backoff *= self.backoff_factor
+                    continue
+
+        raise last_exception
+
+    @contextmanager
+    def create_stream(self, path: Union[str, os.PathLike], mode: str) -> Generator[io.IOBase, None, None]:
+        """Create a stream for reading from or writing to S3 with retry logic."""
+        path_str = str(path)
+        bucket, key = self._parse_s3_uri(path_str)
+        log.info(f"S3 Filesystem: Creating stream for {key} in bucket {bucket}", rank0_only=False)
+
+        if mode == "rb":
+            stream = io.BytesIO()
+            try:
+
+                def download_operation():
+                    self.s3_client.download_fileobj(bucket, key, stream)
+                    stream.seek(0)
+
+                log.info(f"S3 Filesystem: Downloading {key} from bucket {bucket}", rank0_only=False)
+                self._retry_with_backoff(download_operation)
+                log.info("S3 Filesystem: Download complete", rank0_only=False)
+                yield stream
+            finally:
+                stream.close()
+        elif mode == "wb":
+            stream = S3Stream()
+            try:
+                yield stream
+
+                def upload_operation():
+                    stream.seek(0)
+                    self.s3_client.upload_fileobj(stream, bucket, key)
+
+                log.info(f"S3 Filesystem: Uploading {key} to bucket {bucket}", rank0_only=False)
+                self._retry_with_backoff(upload_operation)
+                log.info("S3 Filesystem: Upload complete", rank0_only=False)
+            finally:
+                stream._true_close()
+        else:
+            raise ValueError(f"Unsupported mode: {mode}")
+
+    def concat_path(self, path: Union[str, os.PathLike], suffix: str) -> Union[str, os.PathLike]:
+        """Concatenate S3 path with suffix."""
+        path_str = str(path)
+        if path_str.endswith("/"):
+            return f"{path_str}{suffix}"
+        return f"{path_str}/{suffix}"
+
+    def init_path(self, path: Union[str, os.PathLike]) -> Union[str, os.PathLike]:
+        """Initialize and validate S3 path."""
+        path_str = str(path)
+        if not path_str.startswith("s3://"):
+            raise ValueError(f"Invalid S3 URI: {path_str}. Must start with 's3://'")
+        return path_str
+
+    def rename(self, path: Union[str, os.PathLike], new_path: Union[str, os.PathLike]) -> None:
+        """Rename (move) an object in S3 with retry logic."""
+        src_bucket, src_key = self._parse_s3_uri(str(path))
+        dst_bucket, dst_key = self._parse_s3_uri(str(new_path))
+
+        def copy_operation():
+            copy_source = {"Bucket": src_bucket, "Key": src_key}
+            self.s3_client.copy(copy_source, dst_bucket, dst_key)
+
+        self._retry_with_backoff(copy_operation)
+
+        def delete_operation():
+            self.s3_client.delete_object(Bucket=src_bucket, Key=src_key)
+
+        self._retry_with_backoff(delete_operation)
+
+    def mkdir(self, path: Union[str, os.PathLike]) -> None:
+        """
+        Create a "directory" in S3.
+
+        Note: S3 doesn't have real directories, but we can create an empty object
+        with a trailing slash to simulate a directory.
+        """
+        # Creating same buckets from different ranks can cause rate limit issues in GCP.
+        # In object store, we don't need to create a directory.
+        pass
+
+    @classmethod
+    def validate_checkpoint_id(cls, checkpoint_id: Union[str, os.PathLike]) -> bool:
+        """Validate if the checkpoint_id is a valid S3 URI."""
+        checkpoint_id_str = str(checkpoint_id)
+        try:
+            if not checkpoint_id_str.startswith("s3://"):
+                return False
+            parsed = urlparse(checkpoint_id_str)
+            return bool(parsed.netloc and parsed.path)  # Must have bucket and key
+        except Exception:
+            return False
+
+    def exists(self, path: Union[str, os.PathLike]) -> bool:
+        """Check if an object exists in S3 with retry logic."""
+        bucket, key = self._parse_s3_uri(str(path))
+        try:
+
+            def head_operation():
+                self.s3_client.head_object(Bucket=bucket, Key=key)
+
+            self._retry_with_backoff(head_operation)
+            return True
+        except ClientError as e:
+            if e.response.get("Error", {}).get("Code", "") == "404":
+                return False
+            raise
+
+    def rm_file(self, path: Union[str, os.PathLike]) -> None:
+        """Remove a file from S3 with retry logic."""
+        bucket, key = self._parse_s3_uri(str(path))
+
+        def delete_operation():
+            self.s3_client.delete_object(Bucket=bucket, Key=key)
+
+        self._retry_with_backoff(delete_operation)
+
+    def _parse_s3_uri(self, uri: str) -> tuple[str, str]:
+        """
+        Parse an S3 URI into bucket and key.
+
+        Args:
+            uri: S3 URI in the format s3://bucket-name/key
+
+        Returns:
+            Tuple of (bucket_name, key)
+
+        Raises:
+            ValueError: If the URI is invalid
+        """
+        uri = uri if isinstance(uri, str) else str(uri)
+        if not uri.startswith("s3://"):
+            raise ValueError(f"Invalid S3 URI: {uri}. Must start with 's3://'")
+
+        parsed = urlparse(uri)
+        bucket = parsed.netloc
+
+        # Remove leading slash from key
+        key = parsed.path.lstrip("/")
+
+        if not bucket:
+            raise ValueError(f"Invalid S3 URI: {uri}. No bucket specified")
+
+        return bucket, key
+
+
+class S3StorageWriter(FileSystemWriter):
+    def __init__(
+        self,
+        credential_path: str,
+        path: str,
+        **kwargs,
+    ) -> None:
+        """
+        Initialize an S3 writer for distributed checkpointing.
+
+        Args:
+            region (str): The AWS region for S3.
+            path (str): The S3 URI to write checkpoints to.
+            kwargs (dict): Keyword arguments to pass to the parent :class:`FileSystemWriter`.
+        """
+        super().__init__(
+            path=path,
+            sync_files=False,
+            **kwargs,
+        )
+        self.fs = S3FileSystem(credential_path)  # type: ignore
+        self.path = self.fs.init_path(path)
+
+    @classmethod
+    def validate_checkpoint_id(cls, checkpoint_id: Union[str, os.PathLike]) -> bool:
+        return S3FileSystem.validate_checkpoint_id(checkpoint_id)
+
+
+class S3StorageReader(FileSystemReader):
+    def __init__(self, credential_path: str, path: Union[str, os.PathLike]) -> None:
+        """
+        Initialize an S3 reader for distributed checkpointing.
+
+        Args:
+            region (str): The AWS region for S3.
+            path (Union[str, os.PathLike]): The S3 path to read checkpoints from.
+        """
+        super().__init__(path)
+        self.fs = S3FileSystem(credential_path)  # type: ignore
+        self.path = self.fs.init_path(path)
+        self.sync_files = False
+
+    @classmethod
+    def validate_checkpoint_id(cls, checkpoint_id: Union[str, os.PathLike]) -> bool:
+        return S3FileSystem.validate_checkpoint_id(checkpoint_id)

pid/_ext/imaginaire/config.py

@@ -0,0 +1,450 @@
+# 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.
+
+"""Training config system for Imaginare4"""
+
+from __future__ import annotations
+
+import os
+from typing import Any, Dict, Optional, Type, TypeVar, Union
+
+import attrs
+import torch
+import torch.utils.data
+import torch.utils.data.distributed
+
+try:
+    from megatron.core import ModelParallelConfig
+
+    USE_MEGATRON = True
+except ImportError:
+    USE_MEGATRON = False
+    print("Megatron-core is not installed.")
+
+from pid._ext.imaginaire.lazy_config import LazyCall as L
+from pid._ext.imaginaire.lazy_config import LazyDict
+from pid._ext.imaginaire.utils import callback, distributed
+from pid._ext.imaginaire.utils.misc import Color
+
+T = TypeVar("T")
+
+
+def _is_attrs_instance(obj: object) -> bool:
+    """
+    Helper function to check if an object is an instance of an attrs-defined class.
+
+    Args:
+        obj: The object to check.
+
+    Returns:
+        bool: True if the object is an instance of an attrs-defined class, False otherwise.
+    """
+    return hasattr(obj, "__attrs_attrs__")
+
+
+def make_freezable(cls: T) -> T:
+    """
+    A decorator that adds the capability to freeze instances of an attrs-defined class.
+
+    NOTE: This requires the wrapped attrs to be defined with attrs.define(slots=False) because we need
+    to hack on a "_is_frozen" attribute.
+
+    This decorator enhances an attrs-defined class with the ability to be "frozen" at runtime.
+    Once an instance is frozen, its attributes cannot be changed. It also recursively freezes
+    any attrs-defined objects that are attributes of the class.
+
+    Usage:
+        @make_freezable
+        @attrs.define(slots=False)
+        class MyClass:
+            attribute1: int
+            attribute2: str
+
+        obj = MyClass(1, 'a')
+        obj.freeze()  # Freeze the instance
+        obj.attribute1 = 2  # Raises AttributeError
+
+    Args:
+        cls: The class to be decorated.
+
+    Returns:
+        The decorated class with added freezing capability.
+    """
+
+    if not hasattr(cls, "__dict__"):
+        raise TypeError(
+            "make_freezable cannot be used with classes that do not define __dict__. Make sure that the wrapped "
+            "class was defined with `@attrs.define(slots=False)`"
+        )
+
+    original_setattr = cls.__setattr__
+
+    def setattr_override(self, key, value) -> None:  # noqa: ANN001
+        """
+        Override __setattr__ to allow modifications during initialization
+        and prevent modifications once the instance is frozen.
+        """
+        if hasattr(self, "_is_frozen") and self._is_frozen and key != "_is_frozen":
+            raise AttributeError("Cannot modify frozen instance")
+        original_setattr(self, key, value)  # type: ignore
+
+    cls.__setattr__ = setattr_override  # type: ignore
+
+    def freeze(self: object) -> None:
+        """
+        Freeze the instance and all its attrs-defined attributes.
+        """
+        for _, value in attrs.asdict(self, recurse=False).items():
+            if _is_attrs_instance(value) and hasattr(value, "freeze"):
+                value.freeze()
+        self._is_frozen = True  # type: ignore
+
+    cls.freeze = freeze  # type: ignore
+
+    return cls
+
+
+def _pretty_print_attrs_instance(obj: object, indent: int = 0, use_color: bool = False) -> str:
+    """
+    Recursively pretty prints attrs objects with color.
+    """
+
+    assert attrs.has(obj.__class__)
+
+    lines: list[str] = []
+    for attribute in attrs.fields(obj.__class__):
+        value = getattr(obj, attribute.name)
+        if attrs.has(value.__class__):
+            if use_color:
+                lines.append("   " * indent + Color.cyan("* ") + Color.green(attribute.name) + ":")
+            else:
+                lines.append("   " * indent + "* " + attribute.name + ":")
+            lines.append(_pretty_print_attrs_instance(value, indent + 1, use_color))
+        else:
+            if use_color:
+                lines.append(
+                    "   " * indent + Color.cyan("* ") + Color.green(attribute.name) + ": " + Color.yellow(value)
+                )
+            else:
+                lines.append("   " * indent + "* " + attribute.name + ": " + str(value))
+    return "\n".join(lines)
+
+
+def pretty_print_overrides(overrides: Optional[list[str]] = None, use_color: bool = False) -> str:
+    """
+    Pretty prints overrides.
+    """
+
+    lines: list[str] = []
+    lines.append(Color.cyan("* ") + Color.green("overrides") + ": ")
+    for override in overrides:
+        if override == "--":
+            continue
+        if override.startswith("~"):
+            attribute_name = override[1:]
+            attribute_value = None
+        else:
+            attribute_name, attribute_value = override.split("=")
+        if use_color:
+            lines.append("   " + Color.cyan("* ") + Color.green(attribute_name) + ": " + Color.yellow(attribute_value))
+        else:
+            lines.append("   " + "* " + attribute_name + ": " + str(attribute_value))
+
+    return "\n".join(lines)
+
+
+@make_freezable
+@attrs.define(slots=False)  # slots=False is required for make_freezable. See the make_freezable notes for more info.
+class ObjectStoreConfig:
+    # Whether the file I/O is from object store instead of local disk.
+    enabled: bool = False
+    # Path to the object store credentials file.
+    credentials: str = ""
+    # Object store bucket to read from / write to the objects.
+    bucket: str = ""
+
+
+@make_freezable
+@attrs.define(slots=False)
+class JobConfig:
+    # Project name.
+    project: str = ""
+    # Experiment name.
+    group: str = ""
+    # Run/job name.
+    name: str = ""
+    # W&B mode, can be "online", or "disabled".
+    wandb_mode: str = "online"
+
+    @property
+    def path(self) -> str:
+        return f"{self.project}/{self.group}/{self.name}"
+
+    @property
+    def path_local(self) -> str:
+        local_root = os.environ.get("IMAGINAIRE_OUTPUT_ROOT", "/tmp/imaginaire4-output")
+        return f"{local_root}/{self.path}"
+
+
+@make_freezable
+@attrs.define(slots=False)
+class EMAConfig:
+    # Enable tracking a set of exponential moving average (EMA) weights.
+    enabled: bool = False
+    # EMA decay rate.
+    beta: float = 0.9999
+    # Enable removing "_orig_mod-" from buffer names that is added by torch.compile
+    torch_compile_buffer_renaming: bool = False
+
+
+@make_freezable
+@attrs.define(slots=False)
+class PowerEMAConfig:
+    # Enable tracking a set of exponential moving average (EMA) weights.
+    enabled: bool = False
+    # EDM2 paper EMA decay rate.
+    s: float = 0.1
+    # Enable removing "_orig_mod-" from buffer names that is added by torch.compile
+    torch_compile_buffer_renaming: bool = False
+
+
+@make_freezable
+@attrs.define(slots=False)
+class DDPConfig:
+    # Traverse the computation graph to find parameters that don't receive gradients.
+    find_unused_parameters: bool = False
+    # Set to True if the computation graph does not change during the whole training loop.
+    static_graph: bool = True
+    # Set to True if we want to synchronize buffers. Set to False if the sync is going to be handled elsewhere.
+    broadcast_buffers: bool = True
+
+
+@make_freezable
+@attrs.define(slots=False)
+class CuDNNConfig:
+    # Set to True for better reproducibility of the results (only using deterministic cudnn functions).
+    deterministic: bool = False
+    # If set to True, cudnn will benchmark several algorithms and pick the fastest one.
+    benchmark: bool = True
+
+
+@make_freezable
+@attrs.define(slots=False)
+class JITConfig:
+    # Enable exporting a JIT compiled model.
+    enabled: bool = False
+    # Input tensor shape, for example input.
+    input_shape: Union[list[int], None] = None
+    # Device to compile onto.
+    device: str = "cuda"
+    # # Data type to compile onto.
+    dtype: str = "bfloat16"
+    # Strict mode for PyTorch JIT.
+    strict: bool = True
+
+
+@make_freezable
+@attrs.define(slots=False)
+class CheckpointConfig:
+    # possible checkpoint class
+    type: Optional[Dict] = None
+    # for dcp, whether to use async mode
+    dcp_async_mode_enabled: bool = False
+    # Configs for saving the checkpoints to object store.
+    save_to_object_store: ObjectStoreConfig = attrs.field(factory=ObjectStoreConfig)
+    # Save the checkpoint every N iterations.
+    save_iter: int = 999999999
+    # Configs for loading the checkpoints from object store.
+    load_from_object_store: ObjectStoreConfig = attrs.field(factory=ObjectStoreConfig)
+    # Path of model weights to resume the checkpoint from.
+    load_path: str = ""
+    # Whether to load the training states (optimizer/scheduler/grad-scaler) from the checkpoint path.
+    load_training_state: bool = False
+    # Whether to load the scheduler state only from the checkpoint path. If load_training_state is True, this will be ignored.
+    only_load_scheduler_state: bool = False
+    # Load state_dict to the models in strict mode.
+    strict_resume: bool = True
+    # Configs for JIT compiling EMA model.
+    jit: JITConfig = attrs.field(factory=JITConfig)
+    # Print detailed information during checkpoint saving/loading.
+    verbose: bool = True
+    # keys not to resume from the checkpoint, choices: ["model", "optim", "scheduler", "trainer"]
+    keys_not_to_resume: list[str] = []
+    # Whether to use the local filesystem for broadcasting checkpoint data (used for Tensor Parallel Checkpointer).
+    broadcast_via_filesystem: bool = False
+    load_ema_to_reg: bool = False  # used in inference, load EMA weights to regular model
+    replicate_ema_to_reg_in_training: bool = False  # When loading external ckpt for new-stage training, copy EMA weights to regular model. If self checkpoint, will take no effect.
+    # In dcp planner, skip the weight shape check, load weights into the model even weight shape is different
+    dcp_allow_mismatched_size: bool = False
+
+
+@make_freezable
+@attrs.define(slots=False)
+class NVTXConfig:
+    """Config for NVTX ranges used in the main training loop.
+
+    See tutorials/nanogpt for more details on how to integrate profiling into your model."""
+
+    # Enable the NVTX ranges.
+    enabled: bool = False
+    # Synchronize everything in each NVTX range.
+    cuda_synchronize: bool = False
+
+
+@make_freezable
+@attrs.define(slots=False)
+class StragglerDetectionConfig:
+    """Config for Straggler detection tool: https://gitlab-master.nvidia.com/dl/gwe/fault_tolerance_related/straggler/-/tree/cupti?ref_type=heads"""
+
+    # Enable the Straggler Detection.
+    enabled: bool = False
+    # How frequently should the Straggler reports be generated.
+    report_freq: int = 100
+    # How frequently iterations should be profiled
+    profile_freq: int = 1
+    # What is the maximum relative difference between GPUs after they are considered stragglers
+    max_diff: float = 2.0
+    # Should the error be raised when straggler is detected
+    raise_error: bool = True
+    # Analyze kernels in the forward pass.
+    analyze_forward: bool = True
+    # Analyze kernels in the backward pass.
+    analyze_backward: bool = True
+    # Analyze kernels in the optimizer.
+    analyze_optimizer: bool = True
+    # Analyze dataloading time.
+    analyze_dataloading: bool = True
+
+
+@make_freezable
+@attrs.define(slots=False)
+class Profiling:
+    enable_profiling: bool = False
+    enable_memory_snapshot: bool = False
+    save_s3: bool = False
+    profile_freq: int = 1
+    # Target ranks for profiling, each entry must be >=0 and < world_size.
+    target_ranks: list[int] = list(range(8))
+    # Set `record_shape` and `profile_memory` to False to reduce profile size.
+    record_shape: bool = False
+    profile_memory: bool = False
+    with_stack: bool = True
+    with_modules: bool = True
+
+
+@make_freezable
+@attrs.define(slots=False)
+class TrainerConfig:
+    from pid._ext.imaginaire.trainer import ImaginaireTrainer
+
+    type: Type[ImaginaireTrainer] = ImaginaireTrainer
+    # Set the callback class.
+    # Defaults to the callbacks below.
+    callbacks: LazyDict = LazyDict(
+        dict(
+            ema=L(callback.EMAModelCallback)(),
+            progress_bar=L(callback.ProgressBarCallback)(),
+            wandb=L(callback.WandBCallback)(),
+        )
+    )
+    # distributed parallelism strategy
+    distributed_parallelism: str = "ddp"
+    # Distributed data parallel configs.
+    ddp: DDPConfig = attrs.field(factory=DDPConfig)
+    # cuDNN configs.
+    cudnn: CuDNNConfig = attrs.field(factory=CuDNNConfig)
+    # Set the random seed.
+    seed: int = 0
+    # Set the random seed based on current timestamp
+    timestamp_seed: bool = False
+    # Gradient scaler arguments (for torch.amp.GradScaler).
+    grad_scaler_args: dict = attrs.field(factory=lambda: dict(enabled=False))
+    # Maximum number of iterations to train the model.
+    max_iter: int = 999999999
+    # Maximum number of iterations to validate the model. If None, validate on the entire dataset.
+    max_val_iter: int | None = None
+    # How often we log the training stats.
+    logging_iter: int = 100
+    # Whether we want to run the validation routines.
+    run_validation: bool = True
+    # How often we evaluate on the validation set.
+    validation_iter: int = 999999999
+    # Kill the process after N seconds since the last iteration (usually means dead job).
+    timeout_period: int = 999999999
+    # Tensor memory organization format.
+    memory_format: torch.memory_format = torch.preserve_format
+    # Gradient accumulation (update step every N iteration).
+    grad_accum_iter: int = 1
+    # Straggler Detection config
+    straggler_detection: StragglerDetectionConfig = attrs.field(factory=StragglerDetectionConfig)
+    # Profiling config
+    profiling: Profiling = attrs.field(factory=Profiling)
+
+
+@make_freezable
+@attrs.define(slots=False)
+class Config:
+    """Config for an imaginaire4 job.
+
+    See /README.md/Configuration System for more info.
+    """
+
+    # Model configs.
+    model: LazyDict
+    # Optimizer configs.
+    optimizer: LazyDict
+    # Scheduler configs.
+    scheduler: LazyDict
+    # Training data configs.
+    dataloader_train: LazyDict
+    # Validation data configs.
+    dataloader_val: LazyDict
+
+    # Training job configs.
+    job: JobConfig = attrs.field(factory=JobConfig)
+
+    # Trainer configs.
+    trainer: TrainerConfig = attrs.field(factory=TrainerConfig)
+
+    if USE_MEGATRON:
+        # Megatron-Core configs
+        model_parallel: ModelParallelConfig = attrs.field(factory=ModelParallelConfig)
+    else:
+        model_parallel: None = None
+
+    # Checkpointer configs.
+    checkpoint: CheckpointConfig = attrs.field(factory=CheckpointConfig)
+
+    # enable upload reproducible setup to s3
+    upload_reproducible_setup: bool = False
+
+    def pretty_print(self, use_color: bool = False) -> str:
+        return _pretty_print_attrs_instance(self, 0, use_color)
+
+    def to_dict(self) -> dict[str, Any]:
+        return attrs.asdict(self)
+
+    def validate(self) -> None:
+        """Validate that the config has all required fields."""
+
+        # broadcast job.name across all ranks to make sure it is consistent
+        # otherwise, unaligned job names leads unaligned path to save checkpoints
+        job_name_tensor = torch.ByteTensor(bytearray(self.job.name, "utf-8")).cuda()
+        distributed.broadcast(job_name_tensor, 0)
+        self.job.name = job_name_tensor.cpu().numpy().tobytes().decode("utf-8")
+
+        assert self.job.project != ""
+        assert self.job.group != ""
+        assert self.job.name != ""

pid/_ext/imaginaire/flags.py

@@ -0,0 +1,52 @@
+# 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.
+
+"""Feature flags."""
+
+import os
+from dataclasses import dataclass
+
+
+def _parse_bool(value: str) -> bool:
+    """Parse string to a boolean."""
+    return value.lower() in ["true", "1", "yes", "y"]
+
+
+INTERNAL = _parse_bool(os.environ.get("COSMOS_INTERNAL", "0"))
+"""Whether to enable internal (nvidia-only) features."""
+
+SMOKE = _parse_bool(os.environ.get("COSMOS_SMOKE", "0"))
+"""Whether to enable smoke test.
+
+Disables expensive operations such as checkpoint loading.
+"""
+
+VERBOSE = _parse_bool(os.environ.get("COSMOS_VERBOSE", "0"))
+"""Whether to enable verbose output."""
+
+EXPERIMENTAL_CHECKPOINTS = _parse_bool(os.environ.get("COSMOS_EXPERIMENTAL_CHECKPOINTS", "0"))
+"""Whether to enable experimental checkpoints."""
+
+
+@dataclass
+class Flags:
+    internal: bool = INTERNAL
+    smoke: bool = SMOKE
+    verbose: bool = VERBOSE
+    experimental_checkpoints: bool = EXPERIMENTAL_CHECKPOINTS
+
+
+FLAGS = Flags()
+"""Convenience object for accessing flags."""

pid/_ext/imaginaire/lazy_config/__init__.py

@@ -0,0 +1,60 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+import os
+
+from omegaconf import DictConfig, OmegaConf
+
+from pid._ext.imaginaire.lazy_config.instantiate import instantiate
+from pid._ext.imaginaire.lazy_config.lazy import LazyCall, LazyConfig
+from pid._ext.imaginaire.lazy_config.omegaconf_patch import to_object
+
+OmegaConf.to_object = to_object
+
+PLACEHOLDER = None
+LazyDict = DictConfig
+
+__all__ = ["instantiate", "LazyCall", "LazyConfig", "PLACEHOLDER", "LazyDict"]
+
+
+DOC_BUILDING = os.getenv("_DOC_BUILDING", False)  # set in docs/conf.py
+
+
+def fixup_module_metadata(module_name, namespace, keys=None):
+    """
+    Fix the __qualname__ of module members to be their exported api name, so
+    when they are referenced in docs, sphinx can find them. Reference:
+    https://github.com/python-trio/trio/blob/6754c74eacfad9cc5c92d5c24727a2f3b620624e/trio/_util.py#L216-L241
+    """
+    if not DOC_BUILDING:
+        return
+    seen_ids = set()
+
+    def fix_one(qualname, name, obj):
+        # avoid infinite recursion (relevant when using
+        # typing.Generic, for example)
+        if id(obj) in seen_ids:
+            return
+        seen_ids.add(id(obj))
+
+        mod = getattr(obj, "__module__", None)
+        if mod is not None and (mod.startswith(module_name) or mod.startswith("fvcore.")):
+            obj.__module__ = module_name
+            # Modules, unlike everything else in Python, put fully-qualitied
+            # names into their __name__ attribute. We check for "." to avoid
+            # rewriting these.
+            if hasattr(obj, "__name__") and "." not in obj.__name__:
+                obj.__name__ = name
+                obj.__qualname__ = qualname
+            if isinstance(obj, type):
+                for attr_name, attr_value in obj.__dict__.items():
+                    fix_one(objname + "." + attr_name, attr_name, attr_value)
+
+    if keys is None:
+        keys = namespace.keys()
+    for objname in keys:
+        if not objname.startswith("_"):
+            obj = namespace[objname]
+            fix_one(objname, objname, obj)
+
+
+fixup_module_metadata(__name__, globals(), __all__)
+del fixup_module_metadata

pid/_ext/imaginaire/lazy_config/file_io.py

@@ -0,0 +1,24 @@
+# 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.
+
+from iopath.common.file_io import HTTPURLHandler, OneDrivePathHandler, PathHandler
+from iopath.common.file_io import PathManager as PathManagerBase
+
+__all__ = ["PathManager", "PathHandler"]
+
+
+PathManager = PathManagerBase()
+PathManager.register_handler(HTTPURLHandler())
+PathManager.register_handler(OneDrivePathHandler())