Skip to content

feat(training): decouple frozen Nano Reasoner for offline and remote SFT conditioning - #270

Open
Shn0930 wants to merge 4 commits into
NVIDIA:mainfrom
Shn0930:feat/nano-sft-remote-reasoner
Open

Shn0930 wants to merge 4 commits into
NVIDIA:mainfrom
Shn0930:feat/nano-sft-remote-reasoner

Conversation

@Shn0930

@Shn0930 Shn0930 commented Sep 22, 2026

Copy link
Copy Markdown

Motivation

vision_sft_nano trains only the Generator-side moe_gen, time_embedder,
vae2llm, and llm2vae parameters. The Reasoner remains frozen, but the
original joint path still constructs, FSDP-shards, all-gathers, and executes its
8.19B parameters, while also duplicating them in the FP32 EMA model.

This PR moves the frozen Reasoner behind a provider-neutral per-layer K/V
boundary so Generator training ranks can avoid materializing it.

What changed

Shared conditioning interface

  • Adds joint, inline, offline, and remote Reasoner-conditioning backends.
  • Keeps joint as the default, so existing recipes remain unchanged.
  • Uses exact RoPE-applied Reasoner K/V tensors while Generator Q and Generator
    K/V remain live and trainable.
  • Adds strict shape, dtype, offset, identity, and fingerprint validation.
  • Structurally prunes the UND/Reasoner pathway before materialization and FSDP
    wrapping for external backends, including the EMA model.
  • Preserves Generator parameter FQNs for checkpoint compatibility.

inline is the local two-pass numerical reference. read_through is reserved
in the configuration contract but deliberately remains unimplemented.

Offline backend

  • Adds a strict Reasoner-only DCP loader that does not construct the Generator,
    VAE, or EMA model.
  • Adds deterministic SFT document enumeration using the training framing logic.
  • Adds a resumable distributed extraction CLI.
  • Stores immutable, layer-major K/V data in sharded safetensors files.
  • Adds atomic publication, SHA-256 shard validation, per-record fingerprints,
    and distributed completion checks.
  • Adds generator-only checkpoint warm-start and EMA initialization handling.

Remote backend

  • Adds an optional gRPC transport and a one-process/one-GPU Reasoner service.
  • Adds a versioned startup handshake and feature-signature validation.
  • Streams chunked K/V tensors with per-chunk and whole-response checksums.
  • Adds asynchronous requests, absolute deadlines, bounded transient retries,
    token-based admission, queue limits, and backpressure.
  • Synchronizes startup and request failures across Generator ranks before the
    corresponding FSDP collective.
  • Marks a replica unhealthy after OOM or runtime-invariant violations.
  • Ensures provider resources are released on success, dry-run, and failure.

Install the optional dependencies with:

uv sync --extra train --extra reasoner-remote

Validation

Automated validation

  • 186 targeted regression tests passed.
  • Ruff check and formatting passed.
  • Pyrefly reported 0 errors.
  • git diff --check, pre-commit, and uv lock --check passed.

The tests cover cache integrity, tensor codecs, protocol corruption, retry and
deadline handling, cancellation, admission, OOM propagation, distributed
failure synchronization, checkpoint loading, and resource cleanup.

Offline 8xH20 A/B

Backend Mean step time Mean peak allocated / Generator GPU
Joint Reasoner + Generator 57.887 s 52.128 GiB
Offline K/V + Generator only 48.015 s 36.498 GiB

Across six paired steady-state steps, offline conditioning reduced mean step
time by 17.05%, increased aggregate token throughput by 20.56%, and reduced
peak allocated memory by 15.63 GiB per GPU. This is a short hot-cache benchmark,
not a full-corpus storage benchmark.

Remote correctness and training

A real 36-layer Nano Reasoner on H20 produced bitwise-identical K/V tensors
through direct extraction and localhost gRPC for every layer.

A separate run reserved GPU 0 for the Reasoner service and used GPUs 1--7 for
seven-rank Generator FSDP. The first nonzero-LR update changed all 405 live
Generator tensors. A same-job restart restored model, optimizer, scheduler, and
trainer state at iteration 2, continued at the restored LR, and changed all 405
tensors again at iteration 3.

Peak physical memory was approximately 18.6 GiB on the Reasoner GPU and
30.0--33.0 GiB per Generator GPU. No CUDA OOM, RPC, NCCL, or training error
occurred. Checkpoints contain 405 live plus 405 EMA Generator leaves and no
Reasoner parameters.

Limitations and remaining gates

  • Production remote throughput under representative 1/2/4/8-rank concurrency
    and networking remains unmeasured.
  • Remote loss, gradients, and post-update weights still need an identically
    configured joint/inline/offline comparison.
  • Exact dataloader-position resume and uninterrupted-versus-resumed numerical
    parity remain pending.
  • Representative and maximum-length prompt validation remains pending.
  • Multi-document and per-view-caption inputs, including 7-view/11-view cached
    attention, currently fail closed.
  • read_through, dynamic batching, layerwise H2D staging, automatic artifact
    digest derivation, and persistent service caching are not implemented.
  • TLS and authentication are not implemented; non-loopback serving is for an
    isolated trusted network only.
  • Full-corpus extraction and compile-enabled training benchmarks remain pending.

Raw benchmark logs and 104 GiB checkpoints are local artifacts under outputs/
and are not included in this PR.

Rollout guidance

  • Existing users require no migration because joint remains the default.
  • Prefer offline for a fixed and enumerable SFT corpus.
  • Use remote when prompts or augmentations make exhaustive caching impractical.
  • Keep Reasoner workers outside the Generator FSDP process group.
  • On an eight-GPU node, reserving one GPU for the service leaves seven Generator
    ranks; a separate shared Reasoner pool is preferable for production.
  • Launch distributed training through the module entrypoint:
torchrun --nproc_per_node=<N> --module cosmos_framework.scripts.train ...

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants