small bugs of megatron-fsdp
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 3h
- Merged PRs (30d)
- 272
Description
# Describe the bug
hello, I am reading the code of megatron-fsdp, here are some code snippets I found may have some small issues:
- [line 311 in `megatron_fsdp/param_and_grad_buffer.py`](https://github.com/NVIDIA/Megatron-LM/blob/b8e23d5874211a49789bf6a2c5170549f2496d14/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py#L311) may be dead code.
- [line 316 in `megatron_fsdp/param_and_grad_buffer.py`](https://github.com/NVIDIA/Megatron-LM/blob/b8e23d5874211a49789bf6a2c5170549f2496d14/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py#L316) should be `fragment_items = sorted(fragment_items, key=lambda id_item: -id_item[1].numel())`
- [lin 1384 of `_get_parameter_groups` in `megatron_fsdp/param_and_grad_buffer.py`](https://github.com/NVIDIA/Megatron-LM/blob/b8e23d5874211a49789bf6a2c5170549f2496d14/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py#L1384) should be `if any(is_submodule(m, fsdp_module) for fsdp_module in fsdp_modules):`
- [step3 of `_get_parameter_groups function` in `megatron_fsdp/param_and_grad_buffer.py`](https://github.com/NVIDIA/Megatron-LM/blob/b8e23d5874211a49789bf6a2c5170549f2496d14/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py#L1494) have 2 issues:
- `remaining_params` is always empty.
- when the `chunk_size_factor` is changed by the `math.lcm`, the params already in `same_factor_params` may not meet the conditions.
- [line 3160 in `ParamAndGradBuffer.all_reduce_gradients` in `megatron_fsdp/param_and_grad_buffer.py`](https://github.com/NVIDIA/Megatron-LM/blob/b8e23d5874211a49789bf6a2c5170549f2496d14/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py#L3160) should be `if gbuf is None`
- [line 3200 in `GradReducePipeline`](https://github.com/NVIDIA/Megatron-LM/blob/b8e23d5874211a49789bf6a2c5170549f2496d14/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py#L3200) the `check_nan` parameter is set but never used.
- [line 4181 the parameter `all_gather_pipeline`](https://github.com/NVIDIA/Megatron-LM/blob/b8e23d5874211a49789bf6a2c5170549f2496d14/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py#L4181) seems useless.
If I am wrong, please correct me. Hope these issues can be fixed, thanks!
Contributor guide
Assessment
This issue has not been assessed yet.