NVIDIA / NVIDIA/Megatron-LM

Attribute Error on Saving FSDP_DTensor with EP8

Open
#2,035 1 comment 0 reactions 2 assignees Claimed by @xuwchen View on GitHub
bug community-request module: moe
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

**Describe the bug**

A clear and concise description of what the bug is.

Trying to save DeepSeek style FSDP model:
```
888: [rank888]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/training.py", line 1742, in save_checkpoint_and_time
888: [rank888]: save_checkpoint(
888: [rank888]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/checkpointing.py", line 557, in save_checkpoint
888: [rank888]: state_dict = preprocess_fsdp_dtensor_state_dict(args, state_dict, model[0])
888: [rank888]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
888: [rank888]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/checkpointing.py", line 853, in preprocess_fsdp_dtensor_state_dict
888: [rank888]: model_state_dict, optimizer_state_dict = handle_swiglu_in_state_dict(
888: [rank888]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
168: fs-mbz-gpu-129:2477589:2490929 [0] NCCL INFO Init timings - ncclCommInitRankConfig: rank 0 nranks 1 total 4.05 (kernels 0.00, alloc 0.02, bootstrap 0.00, allgathers 0.09, topo 0.48, graphs 0.00, connections 3.45, rest 0.00)
888: [rank888]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/core/transformer/fsdp_dtensor_checkpoint.py", line 275, in handle_swiglu_in_state_dict
888: [rank888]: weight_w, weight_v = split_swiglu_linear_fc1(
888: [rank888]: ^^^^^^^^^^^^^^^^^^^^^^^^
888: [rank888]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/core/transformer/fsdp_dtensor_checkpoint.py", line 219, in split_swiglu_linear_fc1
888: [rank888]: fsdp_slice = dist_param.megatron_fsdp_slice
888: [rank888]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
888: [rank888]: AttributeError: 'Parameter' object has no attribute 'megatron_fsdp_slice'
```

**Steps/Code to reproduce bug**

Please list *minimal* steps or code snippet for us to be able to reproduce the bug.

A helpful guide on on how to craft a minimal bug report http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports.

**Expected behavior**

A clear and concise description of what you expected to happen.

**Additional context**

Add any other context about the problem here.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.