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

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.

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

Description

There appears to be an argument-passing bug in FSDPLoadPlanner.__init__.

In:

# 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:

# bytecheckpoint/planner/default_planner.py

def __init__(
    self,
    flatten_state_dict: bool = True,
    flatten_sharded_tensors: bool = True,
    strict: bool = True,
) -> None:

Therefore:

FSDPLoadPlanner(False)

is effectively interpreted as:

DefaultLoadPlanner(
    flatten_state_dict=False,
    flatten_sharded_tensors=True,
    strict=True,
)

instead of the expected:

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:

planner = FSDPLoadPlanner(False)

print(planner.strict)
print(planner.flatten_state_dict)

Current behavior:

True
False

Expected behavior:

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:

class FSDPLoadPlanner(DefaultLoadPlanner):
    def __init__(self, strict: bool):
        super().__init__(strict=strict)

This preserves the default values of:

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.