rename nvfsdp to mfsdp globally (#1137)

Rename folders and files to no longer use the deprecated nvFSDP name,
now uses 'megatron-fsdp' or 'mfsdp'

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- New Features/Improvements
- Migrate project references and workflows from NVFSDP to Megatron-FSDP
(mfsdp); W&B run names and recipe paths updated; simplified setup for
one recipe by removing an external tarball download.

- Refactor
- Renamed config/API flags from use_nvfsdp → use_mfsdp and updated
checkpointing and save/load flows to the mfsdp backend.

- Documentation
- All NVFSDP mentions, examples, and readmes replaced with Megatron-FSDP
terminology.

- Tests
  - Test names, configs, and orchestration updated to use mfsdp.

- Chores
- Re-enabled linting for a previously excluded directory; disabled a
small masking-ratio assertion in one test.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Peter St. John <pstjohn@nvidia.com>

Author: pstjohn

bionemo-recipes.md

@@ -8,7 +8,7 @@ The biological AI community is actively prototyping model architectures and need
 
 - **Flexible scaling**: Scale from single-GPU prototyping to multi-node training without complex parallelism configurations
 - **Framework compatibility**: Works with popular frameworks like HuggingFace Accelerate, PyTorch Lightning, and vanilla PyTorch
-- **Performance optimization**: Leverages TransformerEngine and nvFSDP for state-of-the-art training efficiency
+- **Performance optimization**: Leverages TransformerEngine and megatron-fsdp for state-of-the-art training efficiency
 - **Research-friendly**: Hackable, readable code that researchers can easily adapt for their experiments
 
 ### Use Cases
@@ -35,7 +35,7 @@ Example models include ESM-2, Geneformer, and AMPLIFY.
 Self-contained training examples demonstrating best practices for scaling biological foundation models. Each recipe is a complete Docker container with:
 
 - **Framework examples**: Vanilla PyTorch, HuggingFace Accelerate, PyTorch Lightning
-- **Feature demonstrations**: FP8 training, nvFSDP, context parallelism, sequence packing
+- **Feature demonstrations**: FP8 training, megatron-fsdp, context parallelism, sequence packing
 - **Scaling strategies**: Single-GPU to multi-node training patterns
 - **Benchmarked performance**: Validated throughput and convergence metrics
 
@@ -57,7 +57,7 @@ tokenizer = AutoTokenizer.from_pretrained("nvidia/AMPLIFY_120M")
 
 ```bash
 # Navigate to a recipe
-cd recipes/esm2_native_te_nvfsdp
+cd recipes/esm2_native_te_mfsdp
 
 # Build and run
 docker build -t esm2_recipe .
@@ -191,4 +191,4 @@ For technical support and questions:
 
 - Check existing issues before opening a new one
 - Review our training recipes for implementation examples
-- Consult the TransformerEngine and nvFSDP documentation for underlying technologies
+- Consult the TransformerEngine and megatron-fsdp documentation for underlying technologies

models/.ruff.toml

@@ -46,7 +46,6 @@ exclude = [
     "dist",
     "node_modules",
     "venv",
-    "packages/nvFSDP/",
 ]
 
 # Ignore import violations in all `__init__.py` files.

recipes/README.md

@@ -2,7 +2,7 @@
 
 This directory contains self-contained training examples that demonstrate best practices for scaling
 biological foundation models using [TransformerEngine](https://github.com/NVIDIA/TransformerEngine)
-and [nvFSDP](https://github.com/NVIDIA-NeMo/nvFSDP). Each recipe is a complete Docker environment with
+and [megatron-fsdp](https://pypi.org/project/megatron-fsdp/). Each recipe is a complete Docker environment with
 benchmarked training scripts that users can learn from and adapt for their own research.
 
 ## Philosophy
@@ -49,7 +49,7 @@ Follow this naming pattern to clearly communicate what your recipe demonstrates:
 
 Examples:
 
-- `esm2_native_te_nvfsdp/` - ESM-2 with vanilla PyTorch, TransformerEngine, and nvFSDP
+- `esm2_native_te_mfsdp/` - ESM-2 with vanilla PyTorch, TransformerEngine, and megatron-fsdp
 - `amplify_accelerate_fp8/` - AMPLIFY with HuggingFace Accelerate and FP8 training
 - `geneformer_lightning_context_parallel/` - Geneformer with PyTorch Lightning and context parallelism
 
@@ -115,16 +115,16 @@ Your `train.py` should be educational and self-explanatory:
 ```python
 #!/usr/bin/env python3
 """
-ESM-2 training with TransformerEngine and nvFSDP.
+ESM-2 training with TransformerEngine and megatron-fsdp.
 
 This script demonstrates how to:
 1. Load and prepare biological sequence data
 2. Initialize ESM-2 with TransformerEngine layers
-3. Configure nvFSDP for memory-efficient multi-GPU training
+3. Configure megatron-fsdp for memory-efficient multi-GPU training
 4. Implement a training loop with proper checkpointing
 
 Key design decisions:
-- We use nvFSDP ZeRO-3 for maximum memory efficiency
+- We use megatron-fsdp ZeRO-3 for maximum memory efficiency
 - TransformerEngine FP8 is enabled for H100+ hardware
 - Context parallelism handles long biological sequences
 """
@@ -197,7 +197,7 @@ optimizer:
 # Distributed training
 distributed:
   backend: nccl
-  nvfsdp:
+  mfsdp:
     enable: true
     sharding_strategy: zero3
 
@@ -242,7 +242,7 @@ training:
   num_train_steps: 100    # Enough steps for stable metrics
 
 wandb:
-  name: "esm2_nvfsdp_benchmark"
+  name: "esm2_mfsdp_benchmark"
   tags: ["L1", "benchmark", "performance"]
 ```
 
@@ -411,7 +411,7 @@ docker run --rm -it --gpus all my_recipe pytest -v .
 
 For reference implementations, examine existing recipes:
 
-- **`esm2_native_te_nvfsdp/`**: Comprehensive example showing vanilla PyTorch with TE and nvFSDP
+- **`esm2_native_te_mfsdp/`**: Comprehensive example showing vanilla PyTorch with TE and megatron-fsdp
 - **`amplify_accelerate_fp8/`**: HuggingFace Accelerate integration with FP8 training
 - **`geneformer_lightning_context_parallel/`**: PyTorch Lightning with context parallelism for long sequences
 

recipes/esm2_native_te_mfsdp/hydra_config/L0_sanity.yaml

@@ -8,7 +8,7 @@ num_train_steps: 250
 
 # WandB config
 wandb_init_args:
-  name: "esm2_t6_8M_UR50D_nvfsdp_sanity"
+  name: "esm2_t6_8M_UR50D_mfsdp_sanity"
   mode: "offline"
 
 # Learning rate scheduler config

recipes/esm2_native_te_mfsdp/hydra_config/L1_15B_perf_test.yaml

@@ -8,7 +8,7 @@ num_train_steps: 500
 
 # WandB config
 wandb_init_args:
-  name: "esm2_t48_15B_UR50D_nvfsdp_L1_perf"
+  name: "esm2_t48_15B_UR50D_mfsdp_L1_perf"
   project: "bionemo-recipes"
 
 # Optimizer config

recipes/esm2_native_te_mfsdp/train_ddp.py

@@ -51,7 +51,7 @@ def is_main_process(self) -> bool:
 
 @hydra.main(config_path="hydra_config", config_name="L0_sanity", version_base="1.2")
 def main(args: DictConfig) -> float | None:
-    """Train ESM-2 with TE layers using nvFSDP.
+    """Train ESM-2 with TE layers using mfsdp.
 
     Model names are valid ESM-2 model sizes, e.g.:
     - "esm2_t6_8M_UR50D"
@@ -63,7 +63,7 @@ def main(args: DictConfig) -> float | None:
     """
     # Initialize distributed training and create a device mesh for FSDP.
     # We have to create a dummy mesh dimension for context parallel and tensor parallel for things
-    # to work correctly with nvFSDP.
+    # to work correctly with mfsdp.
     dist.init_process_group(backend="nccl")
     dist_config = DistributedConfig()
     torch.cuda.set_device(dist_config.local_rank)

recipes/esm2_native_te_mfsdp/train_mfsdp.py

@@ -55,7 +55,7 @@ def is_main_process(self) -> bool:
 
 @hydra.main(config_path="hydra_config", config_name="L0_sanity", version_base="1.2")
 def main(args: DictConfig) -> float | None:
-    """Train ESM-2 with TE layers using nvFSDP.
+    """Train ESM-2 with TE layers using mfsdp.
 
     Model names are valid ESM-2 model sizes, e.g.:
     - "esm2_t6_8M_UR50D"
@@ -67,7 +67,7 @@ def main(args: DictConfig) -> float | None:
     """
     # Initialize distributed training and create a device mesh for FSDP.
     # We have to create a dummy mesh dimension for context parallel and tensor parallel for things
-    # to work correctly with nvFSDP.
+    # to work correctly with mfsdp.
     dist.init_process_group(backend="nccl")
     dist_config = DistributedConfig()
     torch.cuda.set_device(dist_config.local_rank)