HuggingFace Client API

The HuggingFace Client API lets you federate an existing HuggingFace Trainer or TRL SFTTrainer script without manually calling flare.receive() and flare.send(). Import nvflare.client.hf, patch the trainer, and keep the familiar trainer.evaluate() and trainer.train() calls in the training loop.

Use this API when your local training script already uses HuggingFace Trainer style code and you want FLARE to handle the FL task exchange, global-weight loading, local-budget stopping, checkpoint state, and rank-0 communication.

Minimal Client Change

Start from a normal HuggingFace trainer script. After constructing the trainer, call flare.patch(trainer) and run one evaluate/train pair per FL round:

import nvflare.client.hf as flare

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)

flare.patch(trainer)

while flare.is_running():
    trainer.evaluate()
    trainer.train()

nvflare.client.hf re-exports the standard Client API symbols, so the same module can be used for FLModel, get_job_id(), log(), and other Client API calls. Its is_running() implementation is HuggingFace aware and coordinates with the patched trainer state.

What patch() Does

flare.patch(trainer) wraps the trainer methods and registers FLARE callbacks. During each FL round it:

• initializes the FLARE Client API with the global process rank;

• receives the FL task and global model on rank 0;

• broadcasts task metadata and parameters to other distributed ranks;

• loads received parameters into the HuggingFace model;

• limits each local round to the configured local step or epoch budget;

• restores Trainer checkpoint state between rounds when enabled;

• captures evaluation metrics for server-side model selection;

• sends trained parameters and metrics back to FLARE from rank 0.

The training script still owns model construction, dataset loading, TrainingArguments/SFTConfig, optimizers, schedulers, and any normal HuggingFace callbacks.

Patch Options

The most common options are:

restore_state

Whether to restore HuggingFace Trainer state across FL rounds. This is enabled by default so optimizer, scheduler, and global-step state continue across rounds. When enabled, save_only_model=True is not supported.

params_scope

Which model parameters to exchange. "auto" is the default. It uses adapter-only parameters for PEFT PeftModel trainers and full-model parameters otherwise. Use "model" to force full-model exchange or "adapter" to require PEFT adapter exchange.

server_key_prefix

Optional key prefix to strip from server parameters when loading into the trainer, and to add back when sending results. Use this when the server model wrapper has a different state-dict namespace than the trainer model.

local_epochs / local_steps

Optional explicit local training budget. Specify at most one. If neither is set, the wrapper uses TrainingArguments.max_steps when positive; otherwise it converts TrainingArguments.num_train_epochs to optimizer steps from the trainer dataloader.

load_state_dict_strict

Whether incoming parameters must match the local model keyspace strictly. Keep this enabled unless you intentionally expect partial parameter exchange.

stream_metrics

Whether to stream finite scalar HuggingFace logs through the FLARE tracking API on rank 0.

Example:

flare.patch(
    trainer,
    params_scope="auto",
    server_key_prefix="model.",
    local_epochs=1,
    stream_metrics=True,
)

Metrics And Model Selection

Call trainer.evaluate() before trainer.train() when the server should receive validation metrics before local training:

while flare.is_running():
    metrics = trainer.evaluate()
    trainer.train()

The HuggingFace Client API reports the metrics returned by the trainer. Configure the server-side recipe or selector with a higher-is-better metric key. If you do not want the recipe to select a best model, set key_metric="":

from nvflare.app_opt.pt.recipes.fedavg import FedAvgRecipe

recipe = FedAvgRecipe(
    name="hf_sft",
    model=model,
    min_clients=2,
    num_rounds=3,
    train_script="client.py",
    launch_external_process=True,
    key_metric="",
)

If the server sends an evaluation task, trainer.evaluate() sends metrics without running training. If the server requests submit_model, either trainer.evaluate() or trainer.train() can trigger submission of the last completed checkpoint, falling back to the current in-memory model.

Parameter Scope And Key Prefixes

Use params_scope="auto" for most jobs:

• Full-model SFT uses full model parameters.

• PEFT/LoRA jobs using a PeftModel use adapter parameters only.

Use server_key_prefix only when the server-side model and trainer model state-dict keys differ. For example, if a server wrapper stores the base model under self.model, server keys may look like model.transformer... while the local trainer keys omit model.. In that case:

flare.patch(trainer, server_key_prefix="model.")

For PEFT jobs where the server model already exposes adapter-shaped keys, leave the prefix unset:

flare.patch(trainer, params_scope="auto", server_key_prefix=None)

Job Recipe Setup

Use a recipe to package the trainer script and configure the Client API executor. LLM jobs commonly run in an external process so torchrun, accelerate, CUDA, and other training runtime settings stay isolated from the FLARE client job process.

from nvflare.app_opt.pt.recipes.fedavg import FedAvgRecipe
from nvflare.client.config import ExchangeFormat
from nvflare.recipe import SimEnv

recipe = FedAvgRecipe(
    name="hf_sft",
    model={"class_path": "model.CausalLMModel", "args": {"model_name_or_path": "gpt2"}},
    min_clients=2,
    num_rounds=3,
    train_script="client.py",
    train_args="--model_name_or_path gpt2 --local_epochs 1",
    server_expected_format=ExchangeFormat.PYTORCH,
    launch_external_process=True,
    key_metric="",
)

env = SimEnv(num_clients=2)
recipe.execute(env)

Use server_expected_format=ExchangeFormat.PYTORCH to keep PyTorch tensors and preserve tensor dtypes such as bfloat16. Use ExchangeFormat.NUMPY when the server-side workflow expects NumPy arrays; half-precision tensors are cast to float32 before they are sent to a NumPy server.

Distributed Training

For multi-GPU or multi-node HuggingFace training, launch the script with torchrun or another launcher that initializes torch.distributed and sets global RANK/WORLD_SIZE before calling flare.patch(trainer).

All ranks must call the same patched Trainer methods in the same order:

while flare.is_running():
    trainer.evaluate()
    trainer.train()

Rank 0 is the only process that calls the FLARE Client API receive/send path. Other ranks receive task data through torch.distributed broadcast. If one rank calls trainer.evaluate() while another calls trainer.train() for the same task, the wrapper raises an error instead of allowing a distributed deadlock.

Distributed parameter payloads use in-memory object broadcast by default, so parameter delivery does not require shared storage. To opt into automatic file staging above a size threshold, set:

export NVFLARE_HF_PARAMS_EXCHANGE_STRATEGY=auto

You can force file staging with NVFLARE_HF_PARAMS_EXCHANGE_STRATEGY=file. Both file modes stage rank-0 parameters under <output_dir>/_fl_exchange and require all ranks to share TrainingArguments.output_dir. The auto threshold is controlled by NVFLARE_HF_PARAMS_FILE_EXCHANGE_MIN_BYTES. Independently, restore_state=True checkpoint resume requires all ranks to see the same checkpoint paths under output_dir.

Checkpoint Behavior

With restore_state=True:

• the wrapper resumes from the last FLARE-recorded HuggingFace checkpoint for later rounds;

• TrainingArguments.save_total_limit is set to 2 when it is unset, so the current and previous FL checkpoints can be retained;

• load_best_model_at_end=True is rejected because global best-model selection belongs to the FL server;

• save_only_model=True is rejected because optimizer and scheduler state are needed for Trainer resume.

If you pass your own resume_from_checkpoint to trainer.train(), FLARE uses that checkpoint for the train call and applies received global weights in memory instead of mutating the user-provided checkpoint.

The wrapper keeps the last checkpoint path in memory for the patched trainer process. It does not persist FL checkpoint provenance to disk and does not provide transparent recovery of an in-flight training task that the launcher has already reported as failed. In Phase 1, restore_state=True requires a single trainer process lifecycle; explicitly configured launch_once=False is rejected. Use restore_state=False for per-task trainer launches.

restore_state=False creates fresh Trainer-managed optimizer and scheduler instances for each FL round. It rejects Trainers that already have prebuilt optimizer or lr_scheduler instances at patch time because silently discarding those instances would change user-configured training semantics. Use restore_state=True for prebuilt instances, or let Trainer construct them.

By default, FLARE uses an in-memory global-weight override for verified transformers versions and falls back to checkpoint injection otherwise. To override this selection after validating your environment, set NVFLARE_HF_WEIGHT_OVERRIDE_STRATEGY to one of:

auto

Use the verified-version gate. This is the default.

in_memory

Resume from the previous checkpoint and apply the received global weights in memory during on_train_begin. This avoids rewriting model weights to checkpoint storage each round.

checkpoint_injection

Write received global weights into the recorded checkpoint before resume. Use this when the in-memory override is not compatible with the installed HuggingFace stack or distributed backend.

Unsupported Configurations

The first HuggingFace Client API implementation intentionally does not support:

• DeepSpeed;

• FSDP;

• load_best_model_at_end=True;

• save_only_model=True with restore_state=True;

• more than one patched HuggingFace Trainer in the same Python process.

Troubleshooting

Previous HuggingFace FL task ... is still pending

The loop called flare.is_running() again before completing the current task. Call the expected trainer.train() or trainer.evaluate() first.

Divergent HuggingFace Trainer call across ranks

Distributed ranks called different Trainer methods. Make every rank execute the same evaluate/train sequence.

rank > 0, but torch.distributed is not initialized

A nonzero global rank was detected without an initialized distributed process group. Launch distributed jobs with torchrun or remove the stale RANK environment variable for single-process runs.

None of the model parameters matched

The server and trainer state-dict keyspaces differ. Set the correct server_key_prefix or adjust params_scope.

Complete Example

See the full HuggingFace SFT/PEFT example:

• examples/hello-world/hello-huggingface

For the general Client API concepts, see Client API Usage.