ByteDance-Seed / ByteDance-Seed/ByteCheckpoint

Bug: `FSDPLoadPlanner` passes `strict` as the first positional argument to `DefaultLoadPlanner`

Open Beginner friendly
#19 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.