chore: minor cleanup

Author: le1nux

src/modalities/checkpointing/fsdp/fsdp_checkpoint_loading.py

@@ -109,7 +109,7 @@ def __init__(self, global_rank: int, allow_partial_load: bool = False):
 
         Args:
             global_rank (int): The global rank of the process.
-            allow_partial_load (bool, optional): Whether to allow partial loading of the checkpoint. Defaults to True.
+            allow_partial_load (bool, optional): Whether to allow partial loading of the checkpoint. Defaults to False.
         Returns:
             None
         """

src/modalities/checkpointing/stateful/app_state_factory.py

@@ -47,7 +47,7 @@ def get_raw_app_state(
     def get_dcp_checkpointed_app_state_(
         raw_app_state: AppState,
         checkpoint_dir_path: Path,
-        allow_partial_load: bool = True,
+        allow_partial_load: bool = False,
     ) -> AppState:
         """Loads the checkpointed state dict into the raw AppState object
         (i.e., non-checkpoint loaded AppState) in-place.
@@ -56,7 +56,7 @@ def get_dcp_checkpointed_app_state_(
             raw_app_state (AppState): The raw AppState object. Its ``components_to_load`` policy
                 determines which components are restored.
             checkpoint_dir_path (Path): The path to the checkpoint directory.
-            allow_partial_load (bool, optional): Whether to allow partial loading of the checkpoint. Defaults to True.
+            allow_partial_load (bool, optional): Whether to allow partial loading of the checkpoint. Defaults to False.
 
         Raises:
             RuntimeError: Raises an error if the state dict has already been loaded.

src/modalities/config/config.py

@@ -126,11 +126,6 @@ def parse_sharding_strategy_by_name(cls, name: str) -> ShardingStrategy:
         return parse_enum_by_name(name=name, enum_type=ShardingStrategy)
 
 
-# class DCPCheckpointLoadingConfig(BaseModel):
-#     global_rank: Annotated[int, Field(strict=True, ge=0)]
-#     allow_partial_load: bool = True
-
-
 class FSDP1CheckpointSavingConfig(BaseModel):
     checkpoint_path: Path
     global_rank: Annotated[int, Field(strict=True, ge=0)]