support megatron_fsdp#153
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for Megatron's FullyShardedDataParallel (megatron_FSDP) in the unwrap_model utility function. The review feedback identifies critical compatibility issues with older versions of megatron-core. Specifically, importing FullyShardedDataParallel at the top level can cause an ImportError on older versions, and if it is imported defensively as None, including it in the module_instances tuple will trigger a TypeError during type checking. The reviewer has provided actionable code suggestions to handle these cases safely.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| import torch | ||
| from megatron.core import mpu, tensor_parallel | ||
| from megatron.core.distributed import DistributedDataParallel as DDP | ||
| from megatron.core.distributed import FullyShardedDataParallel as megatron_FSDP |
There was a problem hiding this comment.
Importing FullyShardedDataParallel directly at the top level will raise an ImportError on older versions of megatron-core where this class does not exist. Since the repository supports older versions of megatron-core (e.g., version 0.16.0), we should import it defensively within a try-except block.
| from megatron.core.distributed import FullyShardedDataParallel as megatron_FSDP | |
| try: | |
| from megatron.core.distributed import FullyShardedDataParallel as megatron_FSDP | |
| except ImportError: | |
| megatron_FSDP = None |
| if module_instances is None: | ||
| from megatron.core.distributed import TorchFullyShardedDataParallel as torch_FSDP | ||
| module_instances = (DDP, torch_FSDP, Float16Module) | ||
| module_instances = (DDP, torch_FSDP, Float16Module, megatron_FSDP) |
There was a problem hiding this comment.
If megatron_FSDP is None (due to older megatron-core versions), including it in the module_instances tuple will cause isinstance(model, module_instances) to raise a TypeError because None is not a valid class/type. We should conditionally add megatron_FSDP to the tuple only if it is not None.
| module_instances = (DDP, torch_FSDP, Float16Module, megatron_FSDP) | |
| module_instances = [DDP, torch_FSDP, Float16Module] | |
| if megatron_FSDP is not None: | |
| module_instances.append(megatron_FSDP) | |
| module_instances = tuple(module_instances) |
1 participant
No description provided.