Add HuggingFace Trainer Client API

[chesterxgchen]

Summary

• Add nvflare.client.hf.patch() for HuggingFace Trainer / TRL SFTTrainer federated rounds.
• Support train, evaluate, and submit-model tasks with rank-0 Client API ownership, DDP coordination, checkpoint-backed Trainer state restore, budgeted round execution, metrics capture, PEFT adapter scope, and bf16-to-numpy-safe output conversion.
• Add a clean examples/advanced/hf_client_api Qwen SFT/PEFT example using the default flare.patch(trainer) API and a standard prepare_data.py data-prep step, while leaving the legacy examples/advanced/llm_hf example unchanged.
• Add the HF Client API design doc as implemented PR documentation and keep only one 3rdParty/packaging.LICENSE.txt file for the optional packaging dependency.
• Add focused fake and real dependency contract tests, including transformers min/latest and TRL 0.18/latest CI jobs.

Testing

• python -m pytest tests/unit_test/app_opt/hf -q
• python -m pytest tests/unit_test/recipe/fedavg_recipe_test.py -q
• python -m py_compile examples/advanced/hf_client_api/client.py examples/advanced/hf_client_api/model.py examples/advanced/hf_client_api/prepare_data.py examples/advanced/hf_client_api/job.py
• python examples/advanced/hf_client_api/client.py --help
• python examples/advanced/hf_client_api/prepare_data.py --help
• python examples/advanced/hf_client_api/prepare_data.py --n_clients 2 --data_root /private/tmp/nvflare_hf_client_api_qwen_data_prep
• python examples/advanced/hf_client_api/job.py --help
• python examples/advanced/hf_client_api/job.py --export_config --job_dir /private/tmp/nvflare_hf_client_api_qwen_job_prepared --data_root /private/tmp/nvflare_hf_client_api_qwen_data_prep
• Verified exported launcher config uses python3 -u custom/client.py ... rather than an absolute source path under custom/
• git diff --check upstream/main...HEAD
• ./runtest.sh -s
• Earlier in this PR: ./runtest.sh -l

Review Notes

• The new Qwen example is validated locally through compile/help/export checks. The README now follows the standard hello-world example layout. Full simulation requires downloading the Qwen checkpoint and installing HF/TRL example dependencies.
• FedAvg recipe negate_key_metric pass-through was split into separate PR Expose negate_key_metric in FedAvg recipes #4953. This HF PR keeps the Qwen example independent by disabling recipe best-model selection with key_metric="".
• The TRL contract test constructs real CPU SFTTrainer instances and applies flare.patch(), but does not run a real TRL train/send round. The train/send path is covered by real transformers.Trainer and fake-contract tests.
• train_with_evaluation hard-fail policy currently depends on existing Client API config. The new HF example explicitly calls trainer.evaluate() before trainer.train() each round; recipe-level pass-through can be added later if we want every recipe user to get a hard failure when pre-train eval is omitted.

[greptile-apps]

Greptile Summary

This PR introduces nvflare.client.hf.patch(), a HuggingFace Trainer / TRL SFTTrainer adapter for NVFlare federated learning. It wraps trainer.train() and trainer.evaluate() to handle FL task dispatch (train, evaluate, submit_model), DDP coordination with rank-0 ownership, checkpoint-backed state restore, budgeted round execution, PEFT adapter scope, and bf16-safe numpy conversion.

• Core API (nvflare/app_opt/hf/api.py, utils.py, callbacks.py): ~1,400-line state machine (_HFTaskState) managing FL task lifecycle, weight override strategies (in-memory vs. checkpoint injection), DDP error-broadcast patterns via _run_rank_zero_operation / _gather_rank_status, and a params file-exchange path for large models.
• Example (examples/hello-world/hello-huggingface/): Qwen SFT/PEFT example with a standard flare.patch(trainer) loop, data prep, and job configuration script; CI matrix adds transformers min/latest and TRL 0.18/latest jobs.
• Tests (tests/unit_test/app_opt/hf/): Extensive fake-contract and real-dependency tests covering task dispatch, DDP failure broadcast, budget capture, checkpoint injection, and PEFT adapter scope.

Confidence Score: 5/5

Safe to merge; all previously identified DDP hang classes are addressed and the new code is well-structured with thorough test coverage.

The DDP error-broadcast issues flagged in earlier review rounds have been properly fixed: is_running(), _ensure_task(), _capture_budget_if_needed(), and the checkpoint-injection/state-persistence paths all now broadcast rank-0 errors to all ranks before anyone proceeds. The remaining comments are quality suggestions that do not affect correctness in the expected NVFlare deployment environment.

Files Needing Attention: nvflare/app_opt/hf/api.py deserves a second look for the _train_with_evaluation_enabled() and _params_exchange_strategy() default choices; all other files are straightforward.

Important Files Changed

Filename	Overview
nvflare/app_opt/hf/api.py	Core FL-HF task state machine (~1,400 lines). Handles train/evaluate/submit_model dispatch, DDP rank-0 error broadcast, budget capture, and checkpoint weight injection. DDP hang issues from prior review threads are addressed. One minor inconsistency: _train_with_evaluation_enabled() called on all ranks during __init__, but could silently diverge if non-zero ranks fail get_config().
nvflare/app_opt/hf/utils.py	Parameter utilities: scope resolution, extraction, load, checkpoint read/write (safetensors + torch fallback), file exchange staging. torch.load fallback now carries an explanatory comment about the trusted-local-file assumption. Logic is solid.
nvflare/app_opt/hf/callbacks.py	FLCallback (budget boundary + train_begin/end hooks) and FLMetricsCallback (scalar streaming via SummaryWriter). Clean and minimal; ranks correctly guarded for metrics writing.
nvflare/client/hf/init.py	Public re-export surface for nvflare.client.hf. Correctly re-exports patch from nvflare.app_opt.hf and is_running from hf.api. All expected symbols listed in all.
examples/hello-world/hello-huggingface/client.py	Qwen SFT/PEFT example: DDP setup, dataset loading, SFTTrainer construction, and flare.patch() loop. Handles None metrics from evaluate() return on submit_model tasks. Correct rank/local_rank separation.
tests/unit_test/app_opt/hf/conftest.py	autouse fixture snapshots and restores sys.modules for transformers/peft/nvflare.app_opt.hf entries. Works correctly when tests use import_hf_module (which forces fresh re-imports), but does not copy module objects — repeated mutation of already-loaded modules in unusual test orderings could leak state.
.github/workflows/premerge.yml	Adds hf-client-api-tests matrix: transformers min/latest and TRL 0.18/latest CI jobs using uv for dependency install. Shell quoting of version specs is correct.

Sequence Diagram

sequenceDiagram
    participant User as User Script
    participant Rank0 as Rank 0 (HFTaskState)
    participant RankN as Rank N (HFTaskState)
    participant HFTrainer as HF Trainer
    participant FlareAPI as NVFlare Client API

    User->>Rank0: flare.patch(trainer)
    User->>RankN: flare.patch(trainer)
    note over Rank0,RankN: _init_client_api_for_rank, register FLCallback

    loop FL Rounds
        User->>Rank0: flare.is_running()
        Rank0->>FlareAPI: flare_api.is_running()
        FlareAPI-->>Rank0: True/False
        Rank0->>RankN: "broadcast_object(status, src=0)"

        User->>Rank0: trainer.evaluate() wrapped_evaluate
        Rank0->>FlareAPI: flare_api.receive()
        FlareAPI-->>Rank0: "FLModel (task=train)"
        Rank0->>RankN: _broadcast_task_payload(payload)
        Rank0->>Rank0: _load_global_params_once()
        RankN->>RankN: _load_global_params_once()
        Rank0->>HFTrainer: original_evaluate()
        RankN->>HFTrainer: original_evaluate()

        User->>Rank0: trainer.train() wrapped_train
        note over Rank0,RankN: _capture_budget_if_needed (rank0 broadcasts)
        note over Rank0,RankN: checkpoint injection OR in-memory weight override
        Rank0->>HFTrainer: original_train()
        RankN->>HFTrainer: original_train()
        HFTrainer->>Rank0: FLCallback.on_train_end()
        Rank0->>Rank0: _run_rank_zero_operation(send model)
        Rank0->>FlareAPI: flare_api.send(FLModel)
        Rank0->>RankN: "broadcast_object(ok_status, src=0)"
    end

    User->>Rank0: flare.is_running() False
    note over Rank0,RankN: Loop exits

Loading

_{Reviews (25): Last reviewed commit: "Align HF round metric filtering" | Re-trigger Greptile}

[chesterxgchen]

/build

[chesterxgchen]

/build

[chesterxgchen]

/build

[chesterxgchen]

/build

[chesterxgchen]

/build

[codecov-commenter]

Codecov Report

❌ Patch coverage is 84.17317% with 223 lines in your changes missing coverage. Please review.
✅ Project coverage is 63.59%. Comparing base (de202e8) to head (52dcd27).

Files with missing lines	Patch %	Lines
nvflare/app_opt/hf/api.py	86.07%	127 Missing ⚠️
nvflare/app_opt/hf/utils.py	77.80%	89 Missing ⚠️
nvflare/app_opt/hf/callbacks.py	90.14%	7 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #4948      +/-   ##
==========================================
+ Coverage   63.30%   63.59%   +0.29%     
==========================================
  Files        1021     1026       +5     
  Lines      100386   101795    +1409     
==========================================
+ Hits        63549    64741    +1192     
- Misses      36837    37054     +217     

Flag	Coverage Δ
unit-tests	63.59% <84.17%> (+0.29%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[chesterxgchen]

/build

[chesterxgchen]

/build

[chesterxgchen]

/build

[chesterxgchen]

/build

[chesterxgchen]

/build

[YuanTingHsieh]

The persistent checkpoint provenance appears to support exactly one additional lifecycle: restore_state=True with launch_once=False.

With the default launch_once=True, the same _HFTaskState survives all rounds and already remembers last_checkpoint_path in memory. The example uses this default, so fl_state.json is written every round but provides no value on the normal path.

This also does not currently provide crash recovery. If the Trainer subprocess crashes, LauncherExecutor marks the task failed; NVFlare does not transparently launch another Trainer and reconnect it to the in-flight task. A new job would have a different job ID, so its persisted state is intentionally ignored.

Could we narrow Phase 1 to the existing Lightning contract?

• Keep standard HF checkpoints for optimizer/scheduler restoration.
• Keep the checkpoint path in memory.
• Support restore_state=True only for the single-process lifecycle.
• Remove persistent fl_state.json state and the crash/rejoin guarantee.
• Treat cross-process continuity under launch_once=False as a separate feature requiring an end-to-end launcher test and, ideally, consistent Lightning support.

launch_once is serialized into the external Client API configuration, so patch() can reject an explicitly configured launch_once=False; absence of that field can represent the inherently single-process in-process executor.

This retains the useful HF checkpoint mechanism without introducing an untested recovery contract.

[YuanTingHsieh]

Two additional findings:

Blocking: stateless rounds report zero training steps after round 0

At nvflare/app_opt/hf/api.py:594, on_train_begin() clamps the current HF global_step against the previously recorded train_start_global_step.

With restore_state=False, HF resets state.global_step to zero at the beginning of each fresh train() call. The max() retains the previous round's stale value instead. For example, after round 0 ends at step 2, round 1 trains from HF step 0 to 2 while the recorded baseline remains 2, so on_train_end calculates step_delta = 2 - 2 = 0.

This silently causes:

• NUM_STEPS_CURRENT_ROUND=0; FedAvg treats a non-positive value as weight 1.0, producing incorrect aggregation for heterogeneous clients.
• metric_step_offset does not advance, so later rounds reuse the same metric steps.

For stateless mode, on_train_begin should replace the baseline with the fresh HF state.global_step rather than clamping it against the previous value. Please add a real multi-round restore_state=False test asserting positive round weights and monotonically increasing metric steps.

Documentation: weight-override escape hatch is undiscoverable

NVFLARE_HF_WEIGHT_OVERRIDE_STRATEGY is not documented, although the design says users can select the in-memory strategy after validating a newer Transformers version. Users outside the verified range otherwise have no discoverable way to avoid the automatic checkpoint-injection fallback.

Please either document auto, in_memory, and checkpoint_injection (including their risks), or keep this override private and remove the user-facing promise from the design.

[holgerroth]

Suggestion on example placement: let's move examples/advanced/hf_client_api to examples/hello-world/hello-huggingface.

Rationale:

• This example is exactly the minimal entry point for the new API, and it already matches the hello-<framework> convention — the same file set as hello-lightning (client.py, model.py, job.py, data-prep script, requirements.txt), and the README already follows the hello-world layout. With the synthetic 4-row JSONL data and Qwen2.5-0.5B, it is hello-world scale.
• It slots naturally next to hello-pt / hello-lightning / hello-tf as the HuggingFace entry point.
• examples/advanced/llm_hf then stays the advanced HF tier (large-model streaming, quantization filters, multi-node SLURM). As a tracked follow-up, we should migrate llm_hf to flare.patch() so all HF examples use the HF Client API — that migration also doubles as at-scale validation of the new adapter (large payloads, filters, multi-GPU).

Mechanical changes this implies:

• .github/workflows/premerge.yml: the two TRL matrix legs reference examples/advanced/hf_client_api/... paths.
• Docs: in client_api.rst and client_api_usage.rst, list hello-huggingface in the per-framework example list (next to hello-pt/hello-lightning/hello-tf) and keep the llm_hf link under "Advanced Examples" until it is migrated, instead of replacing it; update the "Complete Example" section in hf_client_api.rst and the link in llm_fine_tuning.rst.
• examples/hello-world/README.md: add a "Hello HuggingFace" entry.
• Inside the example, rename the hf_client_api_qwen default workspace/data paths in job.py/prepare_data.py/README (e.g., /tmp/nvflare/hello-huggingface/...) and the directory references in the README.

[pcnudde]

Behavioral findings from pinned-head QA using real Transformers/TRL trainers, real two-process Gloo, and two-site/two-round federated runs.

[chesterxgchen]

/build

[chesterxgchen]

Addressed the P1 DDP receive-path hang in b3bcd92: rank 0 now wraps the full task receive/classification/server-prefix strip path, and for distributed runs broadcasts a task-dispatch failure payload when receive(), _read_task_kind(), or strip_server_key_prefix() fails before normal task broadcast. Nonzero ranks now receive the rank-0 error and raise instead of blocking in _broadcast_object. Single-process behavior is preserved: non-AgentClosed receive protocol errors still raise the original exception. Added fake-distributed rank-0/rank-1 regressions plus a real two-process Gloo rank-0 receive-failure regression. Validation: python -m pytest tests/unit_test/app_opt/hf/task_dispatch_contract_test.py -q; python -m pytest tests/unit_test/app_opt/hf/distributed_contract_test.py -q; python -m pytest tests/unit_test/app_opt/hf -q; ./runtest.sh -s.

[chesterxgchen]

/build

[chesterxgchen]

Addressed the P1 DDP is_running hang in 908abd7: distributed is_running() now catches unexpected rank-0 Client API errors, logs the failure, broadcasts a falsey value, and returns False on all ranks instead of leaving nonzero ranks blocked in broadcast_object_list. Single-process behavior is preserved by re-raising the original exception when world_size <= 1. Added fake-distributed rank-0/rank-1 regressions plus a real two-process Gloo rank-0 is_running-failure regression. Validation: python -m pytest tests/unit_test/app_opt/hf/task_dispatch_contract_test.py -q; python -m pytest tests/unit_test/app_opt/hf/distributed_contract_test.py -q; python -m pytest tests/unit_test/app_opt/hf -q; ./runtest.sh -s.

[chesterxgchen]

/build

[chesterxgchen]

Addressed the latest review findings in f14fee664:

• Distributed is_running() now broadcasts a structured rank-zero failure and raises the same actionable RuntimeError on every rank; unexpected Client API failures are no longer converted into normal loop termination.
• Positional resume_from_checkpoint arguments are recognized and preserved, preventing NVFlare from also injecting the keyword and causing duplicate-argument failures.
• Strict PEFT loading now rejects incomplete adapter payloads before calling set_peft_model_state_dict().
• Main-branch Hello HuggingFace installation docs now install NVFlare from the repository and install the remaining dependencies explicitly; requirements.txt remains pinned to the first compatible unreleased version for post-release use.
• The design now states that all multi-node ranks must share output_dir and that missing staged files fail loudly.
• The real TRL contract now executes a complete CPU train/send round for both full-model and PEFT SFTTrainer paths.

Validation:

• python -m pytest tests/unit_test/app_opt/hf/task_dispatch_contract_test.py tests/unit_test/app_opt/hf/utils_test.py -q (63 passed)
• python -m pytest tests/unit_test/app_opt/hf/distributed_contract_test.py -q (6 passed)
• python -m pytest tests/unit_test/app_opt/hf -q (102 passed, 2 optional-dependency skips)
• /Users/chesterc/nvflare-env13/bin/python -m pytest tests/unit_test/app_opt/hf/real_transformers_contract_test.py tests/unit_test/app_opt/hf/real_trl_contract_test.py -q (6 passed)
• ./runtest.sh -s (passed)

[chesterxgchen]

/build

[chesterxgchen]

/build

[holgerroth]

Nice support for HF!

[YuanTingHsieh]

I took another pass focused on what is actually blocking versus what is architectural simplification.

Blocking correctness issue

_reset_stateless_trainer_task_state() unconditionally sets trainer.optimizer and trainer.lr_scheduler to None for restore_state=False, including before the first FL train round. Hugging Face supports constructing Trainer with prebuilt instances via optimizers=(optimizer, scheduler). The current reset silently discards that user configuration and lets Trainer create its default optimizer/scheduler instead, changing training semantics from round 1.

Please either:

• preserve/recreate the supplied optimizer and scheduler in their initial state each round, or
• reject restore_state=False when the Trainer was constructed with prebuilt optimizer/scheduler instances.

A real regression should construct Trainer with a distinctive optimizer such as SGD plus a custom scheduler and verify that the configured behavior is not silently replaced. HF documents prebuilt optimizer/scheduler support here: https://huggingface.co/docs/transformers/main/optimizers

Non-blocking simplification concerns

• _params_to_exchange_tensors(), write/read/cleanup_params_exchange_file(), the two exchange environment variables, and several barriers/error paths exist only for the shared-file parameter transport. Since Phase 1 supports replicated DDP rather than FSDP/DeepSpeed, consider broadcasting only task metadata as an object and synchronizing the selected model tensors/buffers directly with dist.broadcast(). This removes the shared-output_dir requirement for parameter delivery. PyTorch also recommends tensor collectives rather than object collectives for tensor data: https://docs.pytorch.org/docs/stable/distributed.html
• If file staging is retained, cleanup_params_exchange_file() is necessary; without file staging, it and _params_to_exchange_tensors() can both disappear. I would prefer file staging to be justified by a measured workload and kept as an explicit optimization rather than becoming the automatic path above 64 MiB.
• The rank-aware nvflare.client.hf.is_running() behavior is necessary for DDP, but the separate _running.py forwarding module is not; it can be collapsed into the public namespace. The pending/aborted exceptions are reasonable fail-fast guards because the core Client API's is_running() prefetches and caches a task. A short code comment explaining that would make the otherwise surprising behavior much clearer.

So my merge position is: fix the prebuilt optimizer/scheduler behavior; the transport/module cleanup can be follow-up work if we intentionally accept the present complexity and shared-storage contract.

[chesterxgchen]

Addressed this review in 2ae0d731b:

• patch(..., restore_state=False) now rejects a Trainer that already has a prebuilt optimizer or LR scheduler. The rejection happens before Client API initialization or Trainer mutation and tells users to use restore_state=True or let Trainer create these objects.
• Added fake-contract coverage for optimizer-only and scheduler-only cases.
• Added a real Transformers regression using SGD (momentum=0.75) plus a custom LambdaLR; it verifies the patch is rejected and the exact configured optimizer/scheduler instances remain attached.
• Changed distributed parameter delivery to in-memory object broadcast by default. Shared-file staging is now explicit via NVFLARE_HF_PARAMS_EXCHANGE_STRATEGY=auto (threshold-based) or file (always). Tests cover all three policies and docs identify the shared-storage requirement only for the opt-in file modes.
• Removed nvflare/client/hf/_running.py; the public namespace now aliases the rank-aware implementation directly.
• Added the requested comment explaining that the pending/aborted guards prevent a second core is_running() task prefetch while the cached task is unfinished.

I retained file staging as an explicit optimization rather than replacing it with direct tensor collectives in this PR. A tensor-collective implementation changes dtype/device/key metadata synchronization and should be evaluated separately against large-model workloads.

Validation:

• python -m pytest tests/unit_test/app_opt/hf/api_contract_test.py tests/unit_test/app_opt/hf/client_hf_contract_test.py tests/unit_test/app_opt/hf/task_dispatch_contract_test.py -q (71 passed)
• /Users/chesterc/nvflare-env13/bin/python -m pytest tests/unit_test/app_opt/hf/real_transformers_contract_test.py -q (5 passed)
• python -m pytest tests/unit_test/app_opt/hf -q (105 passed, 2 skipped)
• ./runtest.sh -s (passed)

[chesterxgchen]

/build

[chesterxgchen]

Addressed the generated training-argument quoting finding in b76a289c9. job.py now applies shlex.quote() to both model_name_or_path and the resolved data_root before constructing train_args, so SubprocessLauncher preserves local paths containing whitespace as single arguments. Validation: whitespace-path argument-tokenization smoke passed; python -m py_compile examples/hello-world/hello-huggingface/job.py passed; ./runtest.sh -s passed.

[chesterxgchen]

/build

[chesterxgchen]

Addressed the two remaining quality observations in 52dcd2714:

• Round-result metric extraction now uses the same finite-scalar normalization as streamed metrics. Tensor-like scalars are converted to Python floats; NaN, infinity, strings, and non-scalar tensors are omitted.
• Added a focused regression covering numeric, tensor-like, non-finite, and non-scalar metric values.
• Added an explanatory comment to the post-checkpoint-injection barrier: it ensures all ranks observe the completed checkpoint update before Hugging Face resumes from it.

Validation: python -m pytest tests/unit_test/app_opt/hf/api_contract_test.py tests/unit_test/app_opt/hf/callbacks_contract_test.py -q (26 passed); python -m pytest tests/unit_test/app_opt/hf -q (106 passed, 2 skipped); ./runtest.sh -s (passed).

[chesterxgchen]

/build