fix: make FSDP2 executor work with ABC+Generic model hierarchy

- Wrap each child module individually, skip root (CPython layout
  conflict between ABC+Generic and FSDP2 __class__ assignment)
- Remove manual unshard in clip_grad_norm (DTensor compatible)
- Fix _no_sync to iterate modules() instead of checking root
- Add reshard after unwrap_model
- Guard __init_subclass__ type resolution against dynamic subclasses
- Add fsdp2 to --parallel_mode CLI choices

scripts/tools/train.py

@@ -327,8 +327,8 @@ def parse_args() -> argparse.Namespace:
         "--parallel_mode",
         type=str,
         default="none",
-        choices=["none", "ddp", "fsdp"],
-        help="Parallel training strategy (none, ddp, fsdp).",
+        choices=["none", "ddp", "fsdp", "fsdp2"],
+        help="Parallel training strategy (none, ddp, fsdp, fsdp2).",
     )
     parser.add_argument(
         "--device_type", type=str, default="cuda", help="Device type to use."
@@ -472,7 +472,7 @@ def train(
     ]
     assert os.path.exists(param_path)
     if nprocs > 1 and parallel_mode == "none":
-        raise ValueError("--nprocs > 1 requires --parallel_mode to be 'ddp' or 'fsdp'")
+        raise ValueError("--nprocs > 1 requires --parallel_mode to be 'ddp', 'fsdp', or 'fsdp2'")
 
     # Load config
     config_path = os.path.join(param_path, "config.json")