update README

Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

Author: weifengpy

distributed/FSDP2/README.md

@@ -1,8 +1,26 @@
 ## FSDP2
 To run FSDP2 on transformer model:
 ```
+cd distributed/FSDP2
 torchrun --nproc_per_node 2 train.py
 ```
+* For 1st time, it creates a "checkpoints" folder and save state dicts there
+* For 2nd time, it loads from previous checkpoints
+
+To enable explicit prefetching
+```
+torchrun --nproc_per_node 2 train.py --explicit-prefetch
+```
+
+To enable mixed precision
+```
+torchrun --nproc_per_node 2 train.py --mixed-precision
+```
+
+To showcse DCP API
+```
+torchrun --nproc_per_node 2 train.py --dcp-api
+```
 
 ## Ensure you are running a recent version of PyTorch:
 see https://pytorch.org/get-started/locally/ to install at least 2.5 and ideally a current nightly build.

distributed/FSDP2/train.py

@@ -57,8 +57,6 @@ def main(args):
     fully_shard(model, **fsdp_kwargs)
 
     inspect_model(model)
-    if args.mixed_precision:
-        inspect_mixed_precision(model)
 
     if args.explicit_prefetching:
         set_modules_to_forward_prefetch(model, num_to_forward_prefetch=2)
@@ -70,6 +68,9 @@ def main(args):
         model.reset_parameters()
     else:
         checkpointer.load_model(model)
+    
+    if args.mixed_precision:
+        inspect_mixed_precision(model)
 
     optim = torch.optim.Adam(model.parameters(), lr=1e-2)
     if checkpointer.last_training_time is not None: