ByteDance-Seed / ByteDance-Seed/ByteCheckpoint
Bug: `FSDPLoadPlanner` passes `strict` as the first positional argument to `DefaultLoadPlanner`
- Dominant language
- Python
- Stars
- 290
- Forks
- 22
- PR merge metrics
- No merged PRs in 30d
Description
## Bug: `FSDPLoadPlanner` passes `strict` as the first positional argument to `DefaultLoadPlanner`
### Description
There appears to be an argument-passing bug in `FSDPLoadPlanner.__init__`.
In:
```python
# bytecheckpoint/planner/fsdp/fsdp_planner.py
class FSDPLoadPlanner(DefaultLoadPlanner):
def __init__(self, strict: bool):
super().__init__(strict)
```
`strict` is passed positionally to `DefaultLoadPlanner.__init__`.
However, the signature of `DefaultLoadPlanner.__init__` is:
```python
# bytecheckpoint/planner/default_planner.py
def __init__(
self,
flatten_state_dict: bool = True,
flatten_sharded_tensors: bool = True,
strict: bool = True,
) -> None:
```
Therefore:
```python
FSDPLoadPlanner(False)
```
is effectively interpreted as:
```python
DefaultLoadPlanner(
flatten_state_dict=False,
flatten_sharded_tensors=True,
strict=True,
)
```
instead of the expected:
```python
DefaultLoadPlanner(
flatten_state_dict=True,
flatten_sharded_tensors=True,
strict=False,
)
```
### Impact
This causes two unexpected behaviors:
1. `strict=False` does not actually disable strict loading.
2. `flatten_state_dict` is unintentionally disabled.
For example:
```python
planner = FSDPLoadPlanner(False)
print(planner.strict)
print(planner.flatten_state_dict)
```
Current behavior:
```text
True
False
```
Expected behavior:
```text
False
True
```
In our case, this causes optimizer checkpoint loading to fail because the planner still behaves as `strict=True`, with errors triggered by unmatched state-dict keys such as the top-level `state` entry.
### Suggested Fix
Pass `strict` explicitly as a keyword argument:
```python
class FSDPLoadPlanner(DefaultLoadPlanner):
def __init__(self, strict: bool):
super().__init__(strict=strict)
```
This preserves the default values of:
```python
flatten_state_dict=True
flatten_sharded_tensors=True
```
while correctly forwarding the requested `strict` value.
Please let me know if you would like me to submit a PR for this fix.
Contributor guide
Research direction
Start in bytecheckpoint/planner/fsdp/fsdp_planner.py and inspect FSDPLoadPlanner.__init__ alongside the signature in bytecheckpoint/planner/default_planner.py. Reproduce the FSDPLoadPlanner(False) case and verify that strict follows the requested value while the flattening defaults remain unchanged.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems
- Issue type
- Bug
- Difficulty
- 1/5
- Estimated time
- Under an hour
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 86/100