[HuggingFace] Multi-node client support

[holgerroth]

Fixes # .

Description

Multi-Node Distributed Training Support for NVFlare LLM Fine-tuning
This PR enables multi-node distributed training with NVFlare for LLM fine-tuning across SLURM clusters.

Key Changes:

• Fixed rank vs local_rank distinction for proper multi-node coordination (only global rank 0 communicates with FL server)
• Created wrapper script (run_multinode_training.sh) to handle srun + torchrun coordination across nodes
• Fixed CUDA device mapping to use local_rank (0-7 per node) instead of global rank
• Added WandB integration with conditional enabling via environment variable
• Updated SLURM script for production mode deployment
• Added comprehensive documentation (MULTINODE.md)

Impact:

• Enables training across 2+ nodes with 8 GPUs per node (16+ total GPUs)
• Supports InfiniBand with GPUDirect RDMA for optimal performance
• Maintains single FL client per site while distributing training across nodes

Tested on: 2 nodes × 8 GPUs (16 total GPUs) with SLURM + InfiniBand

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Quick tests passed locally by running ./runtest.sh.
• In-line docstrings updated.
• Documentation updated.

[ZiyueXu77]

Looks good! some minor comments

[holgerroth]

/build

[chesterxgchen]

add few suggestion

[holgerroth]

Added a separate PR for file restructuring in LLM tutorial #3850 @chesterxgchen

[holgerroth]

/build

[greptile-apps]

Greptile Overview

Greptile Summary

This PR enables multi-node distributed training for NVFlare LLM fine-tuning across SLURM clusters. The implementation properly distinguishes between global rank (for FL operations) and local_rank (for CUDA devices), and introduces a wrapper script pattern to coordinate srun + torchrun across nodes.

Critical Issues Found:

• Deadlock bug in client.py:257: The while flare.is_running(): check only executes on rank 0, causing ranks 1-15 to hang indefinitely when the FL session ends. The is_running status must be broadcast to all ranks before the loop condition.
• Race condition in client.py:143: Directory cleanup uses rank == 0 check, which is correct, but the comment/logic should be verified for consistency.
• GPU detection bug in client_wrapper.sh:28: Fails with range notation in CUDA_VISIBLE_DEVICES (e.g., "0-7" gives NGPUS=1 instead of 8).

Other Issues:

• Documentation in MULTINODE.md doesn't mention the deadlock bug, creating a false sense that the implementation is complete.
• Hardcoded NVIDIA-specific values in nvflare.slurm reduce portability.

Architecture is Sound:
The overall design is well thought out - using a wrapper script to handle multi-node coordination, keeping FL client on rank 0 only, and properly using local_rank for GPU mapping. However, the synchronization implementation is incomplete and will cause production failures.

Confidence Score: 0/5

• This PR has a critical deadlock bug that will cause all multi-node training runs to hang indefinitely when FL sessions end
• The deadlock in client.py:257 is a show-stopper - non-rank-0 processes will never know when to exit the training loop because flare.is_running() is only checked on rank 0 without broadcasting the result. This means every multi-node training run will hang with 15 out of 16 processes stuck in an infinite loop. Additionally, the GPU detection bug in the wrapper script could cause silent failures. While the overall architecture is solid, these bugs must be fixed before merge.
• client.py requires immediate attention for the deadlock bug, client_wrapper.sh needs GPU detection fix, and MULTINODE.md should document the synchronization requirements

Important Files Changed

File Analysis

Filename	Score	Overview
examples/advanced/llm_hf/client.py	0/5	Critical deadlock bug: non-rank-0 processes hang indefinitely when FL session ends (line 257), plus directory cleanup race condition (line 143)
examples/advanced/llm_hf/client_wrapper.sh	3/5	GPU detection breaks with range notation in CUDA_VISIBLE_DEVICES (line 28), otherwise solid wrapper implementation
examples/advanced/llm_hf/MULTINODE.md	2/5	Documentation is comprehensive but misleading - describes a broadcast synchronization pattern that isn't actually implemented in client.py
examples/advanced/llm_hf/nvflare.slurm	4/5	Hardcoded NVIDIA-specific account/partition names need placeholders for portability, otherwise solid SLURM configuration

Sequence Diagram

sequenceDiagram
    participant SLURM as SLURM Master Node
    participant Server as NVFlare Server
    participant Client as NVFlare Client (Rank 0)
    participant Wrapper as client_wrapper.sh
    participant Node0 as Node 0 (Ranks 0-7)
    participant Node1 as Node 1 (Ranks 8-15)
    
    SLURM->>Server: Start NVFlare server
    SLURM->>Client: Start NVFlare client
    SLURM->>Client: Submit FL job via job.py
    
    loop Each FL Round
        Client->>Wrapper: Execute: bash client_wrapper.sh client.py
        Wrapper->>Wrapper: Detect multi-node setup (NNODES=2)
        Wrapper->>Node0: srun launches torchrun --node_rank=0
        Wrapper->>Node1: srun launches torchrun --node_rank=1
        
        Node0->>Node0: Spawn 8 processes (ranks 0-7)
        Node1->>Node1: Spawn 8 processes (ranks 8-15)
        
        Note over Node0,Node1: Only Rank 0 calls flare.receive()
        Server->>Client: Send global model
        Client->>Node0: Rank 0 receives model
        Node0->>Node0: Rank 0 broadcasts model to ranks 1-7
        Node0->>Node1: Rank 0 broadcasts model to ranks 8-15
        
        Note over Node0,Node1: All 16 ranks train via PyTorch DDP
        Node0->>Node0: Training with NCCL P2P/CUMEM
        Node1->>Node1: Training with NCCL P2P/CUMEM
        Node0->>Node1: Cross-node sync via InfiniBand RDMA
        
        Note over Node0,Node1: Only Rank 0 calls flare.send()
        Node0->>Client: Rank 0 sends trained model
        Client->>Server: Submit model updates
        
        Note over Node0,Node1: CRITICAL BUG: Non-rank-0 processes<br/>don't know when to exit loop<br/>(flare.is_running() not broadcast)
    end

Loading

[greptile-apps]

Additional Comments (1)

1. examples/advanced/llm_hf/client.py, line 164-171 (link)

   style: dataset info printing and logging_steps calculation only happen on local_rank == 0, but should use rank == 0 for multi-node consistency (otherwise both nodes print)

_{10 files reviewed, 4 comments}

_{Edit Code Review Agent Settings | Greptile
React with 👍 or 👎 to share your feedback on this new summary format}

[YuanTingHsieh]

good solution and reference for the users, added some suggestions

[greptile-apps]

_{10 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Additional Comments (1)

1. examples/advanced/llm_hf/client.py, line 257-374 (link)

   logic: critical deadlock: non-rank-0 processes hang forever when FL session ends

   flare.is_running() only returns meaningful values on rank 0 (since flare.init(rank=rank) makes it no-op on other ranks). When the FL server signals completion, only rank 0 exits the loop - all other processes continue waiting at line 260's flare.receive() indefinitely.

   need to broadcast loop continuation signal before checking flare.is_running():

   while True:
       # broadcast whether to continue (must happen before any rank-0-only operations)
       if rank == 0:
           should_continue = flare.is_running()
       else:
           should_continue = None
       
       if dist.is_initialized():
           continue_obj = [should_continue]
           dist.broadcast_object_list(continue_obj, src=0)
           should_continue = continue_obj[0]
       
       if not should_continue:
           break
       
       # rest of existing loop code...
       if rank == 0:
           input_model = flare.receive(timeout=600)
           # ... etc
   

_{10 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{10 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

[holgerroth]

/build

[greptile-apps]

_{10 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Additional Comments (1)

1. examples/advanced/llm_hf/client.py, line 143-148 (link)

   logic: use rank == 0 instead of just rank == 0 check - in multi-node setup, both nodes have local_rank == 0, causing race condition on shared filesystem

_{10 files reviewed, 4 comments}

_{Edit Code Review Agent Settings | Greptile}

[holgerroth]

/build

[YuanTingHsieh]

LGTM!