What this page is
If you already have a centralized PyTorch training script (loss + optimizer + DataLoader + a per-epoch loop), this page is the port guide. Most of your code doesn't change. What does:
- Your model class gets a thin
FLModelWrapperaround it (addsget_weights/set_weightsfor the round-trip). - Your per-epoch loop moves inside a
fl_train_model(payload, ...)function — called once per FL round, receives global weights, returns updated weights. - Your dataset stays exactly where it is. Just zip it into shards before upload.
What you don't write at all: the FedAvg aggregator, the per-worker executor, the persistor that saves the global model — the platform supplies those, and they're identical across every FL job in our catalog.
note
Want to read about the FL plumbing instead of porting? Skip to How FL works on ResonTech for the theory deep-dive (FedAvg, per-round loop, persistor lifecycle, tuning lessons).
Required versions
pytorch >= 2.7.1
Python <= 3.11 (Python 3.11 recommended)How It Works (Theory, if you need)
The platform calls your training function once per FL round. Your existing training logic runs inside fl_train_model() - or whatever name you have - you receive global weights at the start and return updated weights at the end.
Round 0 - Initialization
1. Server starts -> persistor instantiates your model class from config_fed_server.json
└─ e.g. model_def.FLModelWrapper(num_classes=3, pretrained=True)
└─ initial weights come from here (e.g. ImageNet pretrained backbone)
└─ persistor calls model.state_dict() -> these become the global weights
2. Server sends initial weights to all clients
3. Each client's executor calls your fl_train_model(payload, ...)
└─ payload["initial_weights"] = the pretrained state_dict from step 1
└─ your code loads them with model.set_weights(payload["initial_weights"])
└─ trains locally -> returns updated weights
4. Server aggregates all client weights (FedAvg) -> saves new global modelRounds 1..N - Training Loop
1. Server sends latest aggregated weights to all clients
2. Each client loads weights -> trains locally -> returns updated weights
3. Server aggregates -> repeatYour existing training loop (epochs, early stopping, scheduler, checkpointing) all runs inside fl_train_model(). The only difference: receive weights from the server at the start and return updated weights at the end.
How params are passed
config_fed_client.json "args"
↓
GenericClientExecutor.__init__(batch_size=2, img_size=512, ...)
↓
payload["env"]["BATCH_SIZE"] = 2, payload["env"]["IMG_SIZE"] = 512, ...
↓
your fl_train_model(payload, ...)1. Required File Structure in your Workspace (Practice)
warn
The folder names
scripts, configs, requirements, model, shards, files_out and model_out are required exact names — the platform resolves paths by these names. Your job folder must live under jobs/ at the root of your bucket.jobs/
{job-name}/
scripts/ <- EXACT NAME REQUIRED
model_def.py <- YOU WRITE THIS
custom_client_executor.py <- auto-generated at workspace init
custom_persistor.py <- auto-generated at workspace init
model_def_wrapper_template.py <- auto-generated at workspace init to help you write model_def.py (not required)
configs/ <- EXACT NAME REQUIRED
config_fed_client.json <- auto-generated from wizard params
config_fed_server.json <- auto-generated from wizard params
meta.json <- auto-generated from wizard params
requirements/ <- EXACT NAME REQUIRED
requirements.txt <- auto-generated (torch, nvflare, etc.)
shards/ <- EXACT NAME REQUIRED
shard_0.zip <- one .zip per GPU worker
shard_1.zip
...
model/ <- EXACT NAME REQUIRED (optional)
checkpoint.pt <- optional base checkpoint (exactly 1 .pt)
files_out/ <- EXACT NAME REQUIRED (written by platform)
model_out/ <- EXACT NAME REQUIRED (written by platform)note
Auto-generated by the Submit Job wizard when you click Build workspace:
custom_client_executor.py, custom_persistor.py, model_def_wrapper_template.py, config_fed_client.json, config_fed_server.json, meta.json and requirements.txt. If you upload directly to the S3 bucket (e.g. via the Python SDK, the Files page upload button, or boto3), create the folders with the exact names above and place your files inside them.2. What You Need to Write
For ready sample see: git hub repo
2.1 config_fed_client.json - Hyperparams
These args are the params you pass to your fl_train_model(). forwarded as payload["env"] from custom_client_executor.py to your train function:
JSON
{
"executor": {
"path": "custom_client_executor.GenericClientExecutor",
"args": {
"data_root": "/var/tmp/nvflare/data", // DON'T CHANGE
"adapter_module": "model_def", // your python file (without .py)
"train_fn": "fl_train_model", // your entry function name
// ↓ CHANGE THESE to match your model ↓
"local_epochs": 3,
"batch_size": 4,
"learning_rate": 0.001,
"num_classes": 3 // MUST match server config
}
}
}2.2 custom_client_executor.py - worker guard to your model
Whatever you passed to "args" in config_fed_client.json you, should copy to client custom_client_executor.py in `__init__` :
Python
def __init__(
self,
data_root: str = "/var/tmp/nvflare/data",
adapter_module: str = "model_def",
train_fn: str = "fl_train_model",
validate_fn: str = "fl_validate_model",
submit_fn: str = "fl_submit_model",
dataset: Optional[dict] = None,
manifest_file: Optional[str] = "manifest.ndjson",
# generic hyperparams (passed to adapter in env):
# here update params to match yours
# --------------------------------
local_epochs: int = 1,
batch_size: int = 32,
learning_rate: float = 1e-3,
img_size: int = 224,
num_classes: int = 0, # optional override; adapter may infer otherwise
cache_dir_name: str = "_img_cache",
# --------------------------------
out_dir_name: str = "runs", # keep this
):
super().__init__()
self.data_root = data_root
self.adapter_module = adapter_module
self.train_fn = train_fn
self.validate_fn = validate_fn
self.submit_fn = submit_fn
self.dataset_spec = dataset
self.manifest_file = manifest_file
# here update params to match yours
# --------------------------------
self.local_epochs = int(local_epochs)
self.batch_size = int(batch_size)
self.learning_rate = float(learning_rate)
self.img_size = int(img_size)
self.num_classes = int(num_classes) if num_classes else 0
self.cache_dir = os.path.join(self.data_root, cache_dir_name)
# --------------------------------
self.out_root = os.path.join(self.data_root, out_dir_name)
self.device = Noneand to the `payload["env"]` dict in `_run_adapter()` :
YAML
payload = {
"dataset": dataset,
"env": {
"SITE_NAME": site_name,
"JOB_ID": job_id,
"CURRENT_ROUND": current_round,
# here update params to match yours
# --------------------------------
"EPOCHS": self.local_epochs,
"BATCH_SIZE": self.batch_size,
"LR": self.learning_rate,
"IMG_SIZE": self.img_size,
"CACHE_DIR": self.cache_dir,
# optional: keep this aligned with server model shape
"NUM_CLASSES": self.num_classes,
# --------------------------------
},
"initial_weights": global_weights, # may be None in round 0
}That's the only changes to `custom_client_executor.py`.
2.3 model_def.py - Your Entry Point (watch generated model_def_wrapper_template.py for reference)
The client executor dynamically imports your module. You need two things:
A) A model wrapper class - used by the server persistor to initialize and manage weights:
Python
class FLModelWrapper(nn.Module):
def __init__(self, num_classes=3, pretrained=True, **kwargs):
# **kwargs is required — the persistor instantiates your class by passing
# the full "args" dict from config_fed_server.json, which may contain keys
# your __init__ doesn't explicitly declare. Without **kwargs the call will
# raise an unexpected-keyword-argument TypeError at round 0.
super().__init__()
self.model = YourModelBackbone(num_classes=num_classes, ...)
# If pretrained=True, loads pretrained weights (e.g. ImageNet).
# On round 0, persistor calls state_dict() on this instance
# -> those weights become the starting point for ALL clients.
def forward(self, x):
return self.model(x)
def set_weights(self, weights, strict=True):
# Called in your fl_train_model() to load server weights (see below).
# Handle key prefix mismatches here if needed.
...
def get_weights(self):
return {k: v.detach().cpu() for k, v in self.state_dict().items()}
# Called in your fl_train_model() (see below) to extract updated weights
# after local training. Must return CPU tensors.
# The executor converts these to numpy for FL transport.
# Must return CPU tensors.note
Always include
**kwargs in your wrapper's __init__ signature — e.g. def __init__(self, classes, epochs=1, **kwargs). The persistor instantiates your class by unpacking the full args block from config_fed_server.json, which may contain extra keys your constructor doesn't declare. Without **kwargs the job will crash at round 0 with an unexpected keyword argument error.B) An fl_train_model() function - wraps your existing training logic:
Python
def fl_train_model(payload: dict, out_dir: str, logger=None) -> dict:
# 1. Read hyperparams from payload["env"] (passed from config_fed_client -> custom_client_executor.py and now here)
env = payload["env"]
num_classes = int(env["NUM_CLASSES"])
lr = float(env["LR"])
# 2. Create model + load global weights
model = FLModelWrapper()
model.set_weights(payload["initial_weights"])
# 3. Load your dataset shard
data_root = payload["dataset"]["data_root"]
# here dispatch shard same as you dispatch dataset for centralzied training, for example:
images_dir = os.path.join(data_root, "images")
masks_dir = os.path.join(data_root, "labels")
# 4. >>> YOUR EXISTING TRAINING LOGIC GOES HERE <<<
# Call your train(), train_epoch(), etc.
# All your losses, schedulers, augmentations - unchanged.
# 5. Return updated weights + sample count
return {
"weights": model.get_weights(), # executor sends these to server
"samples": len(dataset), # used by FedAvg for weighting
"avg_loss": avg_loss, # optional, for monitoring
}note
Key point: the executor doesn't care what happens inside `fl_train_model()`. Use your `DeepLabLoss`, `GradualUnfreezeScheduler`, `WeightedRandomSampler` — whatever you want. Just load the incoming weights at the start and return updated weights at the end.
2.4 config_fed_server.json - Server Model Init
The server persistor controls initial weights via config_fed_server.json:
YAML
"components": [{
"id": "persistor",
"path": "custom_persistor.CustomPersistor",
"args": {
"model": {
"path": "model_def.FLModelWrapper", // <- your wrapper class
"args": {
//your args which you init model, for example:
"num_classes": 3,
"pretrained": false
}
}
}
}],
"workflows": [
{
"id": "scatter_and_gather",
"path": "nvflare.app_common.workflows.scatter_and_gather.ScatterAndGather",
"args": {
"min_clients": 1,
"num_rounds": 5, // define how many FL rounds you want, anything else can be left unchanged
"start_round": 0,
"wait_time_after_min_received": 10,
"train_task_name": "train",
"train_timeout": 0,
"aggregator_id": "aggregator",
"persistor_id": "persistor",
"shareable_generator_id": "shareable_generator"
}
}
]- Custom checkpoint, source_ckpt_file_full_name - load a specific
.pt.In that case you need to upload it to your S3 bucket and point to its folder in the submission UI, for example 'job_for_test/model'
If you don't need to pass checkpoint, remove "source_ckpt_file_full_name": "/app/checkpoints/GLOBAL_MODEL.pt" row at all
3. What You Don't Change
4. Dataset Format
You split your dataset into one zip per GPU worker; each worker trains on its own shard, so your training code runs unchanged. How to split and what each shard needs: Dataset format & sharding.
note
Key point: don't create unnecessary folders, if your shard originally consist of 2 folders : You should keep your data paths same as you have locally, to ensure that you don't need to refactor the code.
images and labels and your code works with those 2 folders, then DON'T do this :// put each 'images' and 'labels' into 'shard' folder
shard/
images/
labels/
// then you make shard.zip from it and submit.
// worker unpacks shard.zip and start working, but omg, it doesn't know that there's some 'shard' folder,
// it just knows about 'images' and 'labels' ...
// Better: just create .zip from your 'images' and 'labels' directly.