[2.7] Pass-Through: Zero Tensor Copy at CJ for Large-Model Federated Training

[chesterxgchen]

Summary

This PR introduces the pass-through architecture for ClientAPILauncherExecutor, eliminating tensor materialisation at the CJ (Client Job) process when large models are exchanged between the FL server and a subprocess agent.

In large-model federated learning (e.g., 7B–70B LLM fine-tuning), the CJ process today acts as a blind relay that fully deserializes and re-serializes every tensor it receives from the FL server before forwarding to the subprocess. For a 70B float16 model, this consumes ~140 GB of CJ memory and requires two complete network transfers. B1 pass-through removes both costs.

---

Problem Description

NVFlare's multi-hop execution path for launch_external_process=True looks like:

FL Server ──serialize──▶ CJ process ──re-serialize──▶ Subprocess agent

Before this PR

Each tensor in the global model is handled as follows at CJ:

1. Server serializes the model and creates a download transaction (tensor data lives on the server).
2. CJ fully downloads every tensor from the server into its own heap, materialising the complete model in CJ memory.
3. CJ re-serializes the model for the subprocess, creating a new download transaction — the subprocess then downloads from CJ.

For large models, this means:

• CJ peak memory = full model size (potentially 100s of GB).
• Two full network transfers: server → CJ, then CJ → subprocess.
• CJ becomes a throughput bottleneck and an OOM risk for any model that doesn't fit in the CJ process's memory.

This is the reason why workflows are infeasible for large models beyond what the CJ machine can hold.

---

Proposed Architecture and Solution

Pass-Through Architecture

With FOBSContextKey.PASS_THROUGH enabled on CJ's cell FOBS context, the data path becomes:

FL Server ──stream──▶ CJ (LazyDownloadRef only, no tensor data)
                          └──forward ref──▶ Subprocess
                                               └──download──▶ FL Server

CJ holds only lightweight placeholders (< 100 bytes per tensor). The subprocess downloads each tensor directly from the FL server — CJ is never involved in the tensor data path.

Key Components

1. FOBSContextKey.PASS_THROUGH (nvflare/fuel/utils/fobs/__init__.py)

A new context key that signals ViaDownloaderDecomposer to skip the download step and create lazy placeholders instead.

2. LazyDownloadRef (via_downloader.py)

A small sentinel object (four fields: fqcn, ref_id, item_id, dot) created by recompose() in PASS_THROUGH mode. It carries the original FL server's FQCN, batch ref_id, intra-batch item ID, and Datum Object Type — everything the subprocess needs to download the tensor directly.

3. _LazyBatchInfo (via_downloader.py)

A named sentinel stored in fobs_ctx[items_key] during PASS_THROUGH receive. Using a typed class (rather than a plain tuple) makes the PASS_THROUGH branch unambiguous and immune to accidental type collisions with real item dicts.

4. LazyDownloadRefDecomposer (via_downloader.py)

A new auto-registered FOBS decomposer for LazyDownloadRef. When CJ re-serializes a task containing LazyDownloadRef objects:

• decompose() delegates to get_dot_handler(lazy.dot) — the original ViaDownloaderDecomposer subclass (e.g., TensorDecomposer, NumpyArrayDecomposer). That handler's _finalize_lazy_batch post-callback re-emits the original server datum (fqcn + ref_id + DOT) so the subprocess knows exactly where to download from. lazy_dot is appended to the encoding dict for routing on the receive side.
• recompose() uses lazy_dot to look up the handler and delegates to handler.recompose(), which retrieves the real tensor from fobs_ctx[handler.items_key] (populated by process_datum() when the subprocess received the forwarded datum).

The dot (Datum Object Type) field on both LazyDownloadRef and _LazyBatchInfo ensures that numpy arrays stay with NumpyArrayDecomposer and PyTorch tensors stay with TensorDecomposer, preserving type safety through the full pass-through hop.

5. ClientAPILauncherExecutor.initialize() (client_api_launcher_executor.py)

On startup, the executor enables PASS_THROUGH on the engine cell's FOBS context:

cell.core_cell.update_fobs_context({FOBSContextKey.PASS_THROUGH: True})

This single line activates the full B1 architecture for every job that uses launch_external_process=True — including llm_hf and any recipe that calls ScriptRunner(launch_external_process=True).

---

Analysis

Why not intercept at the pipe level?

The pipe (CellPipe) operates on already-serialized bytes. Intercepting at the pipe level would require parsing FOBS binary format, re-writing datum references, and re-assembling the byte stream — fragile and tightly coupled to the wire format.

Intercepting at the FOBS decomposer level is the natural extension point: decomposers already control exactly when and how data is materialised. PASS_THROUGH simply adds a "don't materialise" branch to that existing mechanism.

Why dot (Datum Object Type) on LazyDownloadRef?

The subprocess must know which ViaDownloaderDecomposer subclass owns the downloaded data so it can store it in the correct fobs_ctx[items_key] and route recompose() correctly. The dot field, set when the server originally serialized the tensor, carries this type information through the pass-through hop without any type-switching logic.

Memory profile at CJ

Stage	Before (tensor materialised)	After (B1 pass-through)
CJ receive	Full model size (e.g., 140 GB)	~100 bytes per tensor
CJ forward	Creates new download tx	Re-emits original server datum
Subprocess receive	Downloads from CJ	Downloads directly from FL server

---

Benefits

1. Zero tensor copy at CJ — CJ memory footprint is independent of model size.
2. One network transfer instead of two — tensors travel server → subprocess directly.
3. No CJ OOM risk for large models regardless of CJ machine memory capacity.
4. Transparent to job authors — no changes to job configs, training scripts, or recipe APIs; launch_external_process=True automatically activates B1.
5. Type-safe — dot propagation preserves tensor type (numpy / pytorch) through the hop without any if/elif type switching.

---

Backward Compatibility

• All existing jobs using launch_external_process=True automatically benefit. No config or script changes required.
• Jobs using launch_external_process=False (in-process executor) are completely unaffected — ClientAPILauncherExecutor.initialize() is not called.
• For models smaller than the ViaDownloaderDecomposer streaming threshold (2 MB per array), FOBS uses native (inline) serialization regardless of PASS_THROUGH — behaviour is identical to before.
• LazyDownloadRefDecomposer is auto-registered via the existing register_folder mechanism; no explicit registration call is needed by any caller.

---

Changed Files

File	Change
nvflare/fuel/utils/fobs/__init__.py	Add FOBSContextKey.PASS_THROUGH
nvflare/fuel/utils/fobs/decomposers/via_downloader.py	Add LazyDownloadRef, _LazyBatchInfo, PASS_THROUGH branches in process_datum() / recompose(), LazyDownloadRefDecomposer, _finalize_lazy_batch post-callback
nvflare/app_common/executors/client_api_launcher_executor.py	initialize() enables PASS_THROUGH on engine cell

---

Test Strategy

Unit tests — tests/unit_test/fuel/utils/fobs/test_pass_through.py (22 tests)

Test class	What is verified
TestLazyDownloadRef	Construction, __slots__, per-item distinctness
TestLazyBatchInfo	Construction, __slots__, isinstance reliability vs plain tuple
TestProcessDatumPassThrough	PASS_THROUGH stores _LazyBatchInfo, never calls _download_from_remote_cell; normal mode calls download
TestRecomposePassThrough	Returns LazyDownloadRef with correct fqcn, ref_id, item_id from _LazyBatchInfo
TestDecomposeWithLazyDownloadRef	Returns REF encoding; _finalize_lazy_batch post-CB registered once per batch regardless of item count; emitted datum has correct fqcn/ref_id/DOT
TestNoMemoryAccumulation	_CtxKey.OBJECTS absent after PASS_THROUGH (no download transaction opened); DownloadService._tx_table unchanged; 50-cycle repeat produces no state bleed

End-to-end tests — tests/unit_test/fuel/f3/streaming/test_pass_through_e2e.py (5 tests, real TCP Cells)

Test	What is verified
test_arrays_survive_pass_through_hop	Full round-trip: server → CJ (PASS_THROUGH) → subprocess, arrays arrive bit-exact
test_cj_holds_only_lazy_refs_not_tensor_data	After PASS_THROUGH deserialization, CJ holds only LazyDownloadRef, never np.ndarray
test_cj_creates_no_download_transaction	DownloadService._tx_table is unchanged during PASS_THROUGH + re-serialization
test_forwarded_payload_carries_original_server_ref	Forwarded datum contains original server fqcn and ref_id — subprocess downloads from server, not CJ
test_multiple_array_roundtrip	8-array batch all survive with bit-exact values

Integration test — tests/integration_test/data/jobs/pt_large_model_pass_through/

Full-stack integration test using PTClientAPILauncherExecutor with launch_once=True (the pattern used by llm_hf):

• Model: LargeNet — 3-layer MLP with ~8 MB of float32 parameters, well above the 2 MB ViaDownloaderDecomposer streaming threshold. This forces the real B1 code path (streaming + PASS_THROUGH) rather than the native inline path used by small models.
• Client script: Mirrors llm_hf/client.py structure (while flare.is_running() loop, receive / train / send). Uses CPU-only synthetic data — no dataset download required in CI.
• Added to client_api.yml as "run pt-large-model-pass-through".

🤖 Generated with Claude Code

[greptile-apps]

Greptile Summary

Introduces pass-through architecture that eliminates tensor materialization at the CJ (Client Job) process for large-model federated learning. Instead of downloading and re-serializing model tensors at the CJ (consuming 140+ GB for a 70B model), the CJ now creates lightweight LazyDownloadRef placeholders (~100 bytes each) that preserve the original FL server's FQCN and batch reference. The subprocess agent downloads tensors directly from the FL server, bypassing the CJ entirely.

Key Changes:

• Added FOBSContextKey.PASS_THROUGH flag and LazyDownloadRef/_LazyBatchInfo sentinels to FOBS serialization
• ViaDownloaderDecomposer.process_datum() skips download when PASS_THROUGH=True, storing batch metadata instead
• LazyDownloadRefDecomposer re-emits original server datum during CJ re-serialization via _finalize_lazy_batch callback
• ClientAPILauncherExecutor.initialize() enables PASS_THROUGH on engine cell, with proper restore in finalize()
• Comprehensive test coverage: 22 unit tests, 5 E2E tests with real cells, full integration test with 8 MB model

Impact:

• Zero tensor copy at CJ — memory footprint independent of model size
• One network transfer instead of two (server → subprocess directly)
• Transparent to job authors — automatic for launch_external_process=True
• Type-safe via dot (Datum Object Type) propagation through the hop

Confidence Score: 5/5

• This PR is safe to merge with minimal risk
• The implementation is architecturally sound with excellent test coverage (27 tests total), proper exception handling, backward compatibility guarantees, and clear documentation. The pass-through logic is well-isolated to the FOBS decomposer layer with type-safe sentinel classes. Auto-registration via register_folder ensures the LazyDownloadRefDecomposer is available. The initialize/finalize restoration logic prevents state leakage.
• No files require special attention

Important Files Changed

Filename	Overview
nvflare/fuel/utils/fobs/decomposers/via_downloader.py	Added LazyDownloadRef, _LazyBatchInfo sentinels, LazyDownloadRefDecomposer, and PASS_THROUGH logic to enable zero-copy tensor forwarding at CJ
nvflare/fuel/utils/fobs/init.py	Added FOBSContextKey.PASS_THROUGH constant to enable pass-through mode in FOBS context
nvflare/app_common/executors/client_api_launcher_executor.py	Enable PASS_THROUGH on engine cell in initialize(), restore in finalize(), with proper exception handling
tests/unit_test/fuel/utils/fobs/test_pass_through.py	Comprehensive unit tests (22 tests) covering LazyDownloadRef, _LazyBatchInfo, process_datum, recompose, decompose, and memory accumulation
tests/unit_test/fuel/f3/streaming/test_pass_through_e2e.py	E2E tests (5 tests) with real TCP cells validating round-trip correctness, CJ memory isolation, and download service transaction handling

Sequence Diagram

sequenceDiagram
    participant Server as FL Server
    participant CJ as CJ Process
    participant Sub as Subprocess Agent

    Note over Server,Sub: Before This PR (Tensor Materialization)
    Server->>CJ: serialize tensors
    Note over CJ: Downloads & materializes<br/>full model (e.g., 140 GB)
    CJ->>CJ: Re-serialize for subprocess
    CJ->>Sub: New download transaction
    Sub->>CJ: Download tensors
    Note over CJ: CJ memory = full model size

    Note over Server,Sub: After This PR (Pass-Through)
    Server->>CJ: serialize tensors with ref
    Note over CJ: PASS_THROUGH enabled<br/>Creates LazyDownloadRef<br/>(~100 bytes per tensor)
    CJ->>Sub: Forward LazyDownloadRef<br/>(original server FQCN + ref_id)
    Sub->>Server: Download tensors directly
    Note over CJ: CJ memory ≈ 0 bytes<br/>(no tensor data)

Loading

_{Last reviewed commit: 741b460}

[greptile-apps]

_{11 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[chesterxgchen]

/build

[greptile-apps]

_{13 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[chesterxgchen]

/build

[greptile-apps]

_{14 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{12 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[pcnudde]

Are there any possible incompatibilities with things like filters or other users code that expects real tensors instead of the the lazy references.

[chesterxgchen]

/build

[greptile-apps]

_{14 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[chesterxgchen]

Are there any possible incompatibilities with things like filters or other users code that expects real tensors instead of the the lazy references.
There is no filter between Client Job Process and Client Agent (trainer). No other components.

[chesterxgchen]

/build

[greptile-apps]

_{14 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[chesterxgchen]

/build

[greptile-apps]

_{14 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{14 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[chesterxgchen]

/build