Add FSDP2 functionality

Author: Le

README.md

@@ -173,6 +173,14 @@ To launch single-node multi-gpu training, set the 'multi-gpu' argument
 python -m alf.bin.train --conf=CONF_FILE --root_dir=LOG_DIR --distributed multi-gpu
 ```
 
+DistributedDataParallel (DDP) is used by default. To shard model parameters,
+gradients, and optimizer state with PyTorch FSDP2, add
+`--distributed_strategy fsdp2`:
+```bash
+python -m alf.bin.train --conf=CONF_FILE --root_dir=LOG_DIR \
+    --distributed multi-gpu --distributed_strategy fsdp2
+```
+
 To launch multi-node multi-gpu training, we use torch distributed launch module. The 'local_rank' for each process can be obtained from 'PerProcessContext' class, which can be used to assign gpu for your environment if you wish. For details on how PyTorch assign 'local_rank' and 'ddp_rank', please refer to the [documentation](https://github.com/pytorch/pytorch/blob/main/torch/distributed/launch.py).  To start training, run the following command on the host machine:
 ```bash
 export NCCL_SOCKET_IFNAME=SOCKET # find in ifconfig

alf/algorithms/algorithm.py

@@ -153,6 +153,7 @@ def __init__(self,
         self._replay_buffer = None
 
         self._ddp_activated_rank = -1
+        self._distributed_strategy = 'ddp'
 
         # These 3 parameters are only set when ``set_replay_buffer()`` is called.
         self._replay_buffer_num_envs = None
@@ -333,16 +334,20 @@ def use_rollout_state(self):
         """
         return self._use_rollout_state
 
-    def activate_ddp(self, rank: int):
+    def activate_ddp(self, rank: int, strategy: str = 'ddp'):
         """Prepare the Algorithm with DistributedDataParallel wrapper
 
         Note that Algorithm does not need to remember the rank of the device.
 
         Args:
             rank (int): DDP wrapper needs to know on which GPU device this
                 module's parameters and buffers are supposed to be.
+            strategy: distributed implementation, ``'ddp'`` or ``'fsdp2'``.
         """
+        if strategy not in ('ddp', 'fsdp2'):
+            raise ValueError("Unknown distributed strategy: %s" % strategy)
         self._ddp_activated_rank = rank
+        self._distributed_strategy = strategy
 
     @use_rollout_state.setter
     def use_rollout_state(self, flag):
@@ -951,7 +956,11 @@ def state_dict(self, destination=None, prefix='', visited=None, **kwargs):
         return destination
 
     @common.add_method(nn.Module)
-    def load_state_dict(self, state_dict, strict=True, skip_preloded=True):
+    def load_state_dict(self,
+                        state_dict,
+                        strict=True,
+                        skip_preloded=True,
+                        assign=False):
         """Load state dictionary for the algorithm.
 
         Args:
@@ -964,6 +973,8 @@ def load_state_dict(self, state_dict, strict=True, skip_preloded=True):
             skip_preloded (bool): whether to skip the modules that support
                 pre-loading and have been pre-loaded. Currently only Algorithm
                 and its derivatives support pre-loading. (Default: ``True``)
+            assign (bool): accepted for compatibility with PyTorch's module
+                state-dict APIs. ALF preserves its existing copy semantics.
         Returns:
             namedtuple:
             - missing_keys: a list of str containing the missing keys.

alf/algorithms/rl_algorithm.py

@@ -31,7 +31,8 @@
                                  BasicRLInfo)
 from alf.utils import common, dist_utils, summary_utils
 from alf.utils.summary_utils import record_time
-from alf.utils.distributed import data_distributed_when, make_ddp_performer
+from alf.utils.distributed import (data_distributed_when,
+                                   make_distributed_performer)
 from alf.tensor_specs import TensorSpec
 from .config import TrainerConfig
 
@@ -786,14 +787,21 @@ def _train_iter_on_policy(self):
         return steps
 
     def _unroll(self, unroll_length: int):
+        if self._distributed_strategy == 'fsdp2' and self._ddp_activated_rank != -1:
+            # FSDP parameters are sharded between calls, so every off-policy
+            # unroll must run through the FSDP root's gather/reshard hooks.
+            self._first_unroll = False
+            performer = make_distributed_performer(self, self.unroll.__func__)
+            return performer(unroll_length)
         if self._first_unroll:
             self._first_unroll = False
             if self._ddp_activated_rank != -1:
                 # Even though we don't update parameters during unroll, we still need to
                 # wrap self.unroll in DDP so that the parameters are synchronized across
                 # all workers before the unroll starts. Otherwise, the parameters across
                 # the workers are different for the first unroll.
-                performer = make_ddp_performer(self, self.unroll.__func__)
+                performer = make_distributed_performer(self,
+                                                       self.unroll.__func__)
                 return performer(unroll_length)
 
         return self.unroll(unroll_length)

alf/bin/train.py

@@ -43,7 +43,8 @@
 
 In case you have multiple GPUs on the machine and you would like to
 train with all of them, specify --distributed multi-gpu. This will use
-PyTorch's DistributedDataParallel for training.
+PyTorch's DistributedDataParallel for training. Add
+--distributed_strategy fsdp2 to use composable FullyShardedDataParallel.
 
 If instead of Gin configuration file, you want to use ALF python conf file, then
 replace the "--gin_file" option with "--conf", and "--gin_param" with "--conf_param".
@@ -83,6 +84,9 @@ def _define_flags():
     flags.DEFINE_enum(
         'distributed', 'none', ['none', 'multi-gpu', 'multi-node-multi-gpu'],
         'Set whether and how to run training in distributed mode.')
+    flags.DEFINE_enum(
+        'distributed_strategy', 'ddp', ['ddp', 'fsdp2'],
+        'Parameter distribution strategy used in distributed mode.')
     flags.DEFINE_integer(
         'num_gpus_per_ddp_worker', 1,
         "The number of gpus per DDP worker. If specified will create N DDP workers where each worker "
@@ -201,8 +205,11 @@ def _train(root_dir, local_rank=-1, rank=0, world_size=1):
                 alg_wrapper_ctor = DistributedUnroller
         else:
             alg_wrapper_ctor = None
-        trainer = policy_trainer.RLTrainer(trainer_conf, ddp_rank,
-                                           alg_wrapper_ctor)
+        trainer = policy_trainer.RLTrainer(
+            trainer_conf,
+            ddp_rank,
+            alg_wrapper_ctor,
+            distributed_strategy=FLAGS.distributed_strategy)
     elif trainer_conf.ml_type == 'sl':
         # NOTE: SLTrainer does not support distributed training yet
         if world_size > 1:

alf/trainers/policy_trainer.py

@@ -377,11 +377,14 @@ def train(self):
             self._save_checkpoint()
             checkpoint_saved = True
         finally:
+            fsdp2 = getattr(self._algorithm, '_distributed_strategy',
+                            'ddp') == 'fsdp2'
             if (self._config.save_checkpoint_upon_crash
-                    and not checkpoint_saved and self._rank <= 0):
+                    and not checkpoint_saved and self._rank <= 0
+                    and not fsdp2):
                 self._save_checkpoint()
             elif (self._config.confirm_checkpoint_upon_crash
-                  and not checkpoint_saved and self._rank <= 0):
+                  and not checkpoint_saved and self._rank <= 0 and not fsdp2):
                 # Prompts for checkpoint only when running single process
                 # training (rank is -1) or master process of DDP training (rank
                 # is 0).
@@ -516,12 +519,23 @@ def _save_checkpoint(self):
         # training (rank is -1) or the master process of DDP training (rank is
         # 0). Other DDP ranks only save their local replay buffers.
         global_step = alf.summary.get_global_counter()
+        from alf.utils.distributed import (fsdp2_full_state_dict,
+                                           is_fsdp2_module)
+        fsdp2 = is_fsdp2_module(self._algorithm)
+        algorithm_state = None
+        if fsdp2:
+            # Full FSDP2 state gathering is collective even though only rank 0
+            # writes the resulting checkpoint.
+            algorithm_state = fsdp2_full_state_dict(self._algorithm)
         if self._rank <= 0:
             # Replay buffers are saved separately below as sharded per-rank
             # source files, so rank 0's full checkpoint only contains model,
             # optimizer, metrics, and trainer progress.
-            self._checkpointer.save(global_step=global_step,
-                                    including_replay_buffer=False)
+            self._checkpointer.save(
+                global_step=global_step,
+                including_replay_buffer=False,
+                state_overrides={'algorithm': algorithm_state}
+                if fsdp2 else None)
         # Every rank writes one sharded replay-buffer source file into the
         # checkpoint directory. Restore will redistribute these files across the
         # active worker count, which may differ from the save-time worker count.
@@ -588,7 +602,8 @@ class RLTrainer(Trainer):
     def __init__(self,
                  config: TrainerConfig,
                  ddp_rank: int = -1,
-                 algorithm_wrapper_ctor: Callable = None):
+                 algorithm_wrapper_ctor: Callable = None,
+                 distributed_strategy: str = 'ddp'):
         """
 
         Args:
@@ -599,6 +614,7 @@ def __init__(self,
                 process training.
             algorithm_wrapper_ctor: if not None, will be used to wrap
                 ``self._algorithm_ctor`` before creating ``self._algorithm``.
+            distributed_strategy: ``'ddp'`` or ``'fsdp2'``.
         """
         super().__init__(config, ddp_rank)
 
@@ -669,7 +685,7 @@ def __init__(self,
         self._algorithm.set_path('')
         if ddp_rank >= 0:
             # Activate the DDP training
-            self._algorithm.activate_ddp(ddp_rank)
+            self._algorithm.activate_ddp(ddp_rank, distributed_strategy)
             # Make sure the BN statistics of different processes are synced
             # https://pytorch.org/docs/stable/generated/torch.nn.SyncBatchNorm.html#torch.nn.SyncBatchNorm
             # This conversion needs to be performed before wrapping modules with DDP.
@@ -763,7 +779,24 @@ def _train(self):
             iter_num += 1
             self._trainer_progress.update(iter_num, total_time_steps)
 
-            if self._need_to_evaluate(iter_num):
+            need_to_evaluate = self._need_to_evaluate(iter_num)
+            from alf.utils.distributed import is_fsdp2_module
+            if is_fsdp2_module(self._algorithm):
+                # Only rank 0 owns an evaluator, but FSDP2 parameter gathering
+                # is collective. Broadcast rank 0's decision so every rank
+                # enters the gather/reshard operations in the same order.
+                evaluate_flag = torch.tensor(int(need_to_evaluate),
+                                             device=alf.get_default_device())
+                torch.distributed.broadcast(evaluate_flag, src=0)
+                if bool(evaluate_flag):
+                    performer = self._algorithm._fsdp2_performer
+                    performer.unshard()
+                    if self._rank == 0:
+                        self._eval()
+                        self._num_evals_performed += 1
+                    torch.distributed.barrier()
+                    performer.reshard()
+            elif need_to_evaluate:
                 self._eval()
                 self._num_evals_performed += 1
 
@@ -837,6 +870,11 @@ def _check_dpp_paras_consistency(self, iter_num: int,
         if not training_started:
             return
 
+        # FSDP2 ranks intentionally hold different parameter shards, so the
+        # replicated-parameter consistency check is not applicable.
+        if getattr(self._algorithm, '_distributed_strategy', 'ddp') == 'fsdp2':
+            return
+
         proc_cxt = PerProcessContext()
         if not (proc_cxt.is_distributed
                 and self._config.ddp_paras_check_interval > 0

alf/utils/checkpoint_utils.py

@@ -215,8 +215,18 @@ def _convert_legacy_parameter(checkpoint):
 
         def _load_one(module, checkpoint):
             if isinstance(module, nn.Module):
-                missing_keys, unexpected_keys = module.load_state_dict(
-                    checkpoint, strict=strict)
+                from alf.utils.distributed import (FSDP2_OPTIMIZER_STATE,
+                                                   is_fsdp2_module,
+                                                   load_fsdp2_full_state_dict)
+                fsdp2_checkpoint = (FSDP2_OPTIMIZER_STATE in checkpoint
+                                    or not any('_optimizers.' in key
+                                               for key in checkpoint))
+                if is_fsdp2_module(module) and fsdp2_checkpoint:
+                    missing_keys, unexpected_keys = load_fsdp2_full_state_dict(
+                        module, checkpoint, strict=strict)
+                else:
+                    missing_keys, unexpected_keys = module.load_state_dict(
+                        checkpoint, strict=strict)
             else:
                 module.load_state_dict(checkpoint)
                 missing_keys, unexpected_keys = [], []
@@ -363,7 +373,10 @@ def _separate_state(state):
         replay_buffer_state = {}
 
         for k, v in state.items():
-            if k.find('_optimizers.') >= 0 and isinstance(
+            from alf.utils.distributed import FSDP2_OPTIMIZER_STATE
+            if k == FSDP2_OPTIMIZER_STATE:
+                optimizer_state[k] = v
+            elif k.find('_optimizers.') >= 0 and isinstance(
                     v, dict) and 'param_groups' in v:
                 optimizer_state[k] = v
             elif Checkpointer._is_replay_buffer_key(k):
@@ -718,7 +731,8 @@ def save_replay_buffer(self,
     def save(self,
              global_step,
              suffix: Optional[str] = None,
-             including_replay_buffer=True):
+             including_replay_buffer=True,
+             state_overrides=None):
         """Save states of all modules to checkpoint
 
         Args:
@@ -729,6 +743,9 @@ def save(self,
                 If provided, it will be used as the suffix instead of ``global_step``.
             including_replay_buffer (bool): whether save replay buffer state in
                 the main replay buffer checkpoint file.
+            state_overrides (dict|None): precomputed states keyed by module
+                name. Used by collective state-dict implementations such as
+                FSDP2.
         """
         suffix = suffix or str(global_step)
 
@@ -744,10 +761,12 @@ def save(self,
                         (replay_buffer, is_checkpoint_enabled(replay_buffer)))
                     enable_checkpoint(replay_buffer, False)
         try:
+            state_overrides = state_overrides or {}
             state = {
                 k:
-                    v.module.state_dict()
-                    if type(v) == torch.nn.DataParallel else v.state_dict()
+                    state_overrides[k] if k in state_overrides else
+                    (v.module.state_dict()
+                     if type(v) == torch.nn.DataParallel else v.state_dict())
                 for k, v in self._modules.items()
             }
         finally: