Fix and improve loading of distributed checkpoints

Dockerfile

@@ -30,7 +30,7 @@ ENV PIP_CONSTRAINT=""
 # We need to compile from the repo because of https://github.com/state-spaces/mamba/issues/720 (same for causal-conv1d)
 # We set the number of workers to avoid OOM when compiling on laptop. (TODO: Can we make it configurable?)
 RUN MAX_JOBS=4 pip install --no-build-isolation  "causal-conv1d@git+https://github.com/Dao-AILab/causal-conv1d@2a288a1"
-RUN MAX_JOBS=4 pip install --no-build-isolation "mamba_ssm[causal-conv1d]@git+https://github.com/state-spaces/mamba@74729d0"
+RUN MAX_JOBS=4 pip install --no-build-isolation "mamba_ssm[causal-conv1d]@git+https://github.com/state-spaces/mamba@4a8a2a2"
 # Copy dependency files with universal write permissions for all users.
 COPY --chmod=777 setup.py setup.cfg pyproject.toml ./
 COPY --chmod=777 ./fast_llm/__init__.py fast_llm/

fast_llm/engine/checkpoint/distributed.py

@@ -53,7 +53,7 @@ def load(self, config: CheckpointLoadConfig) -> dict[str, typing.Any] | None:
         loaded_metadata = self._model.config.load_metadata(config.to_copy({"load_config": ModelConfigType.fast_llm}))
         shard_names = self.get_shard_names(config)
         # Make sure all shards to load are in the checkpoint.
-        Assert.leq(set(self.get_shard_names(config)), set(loaded_metadata.shards))
+        Assert.leq(set(shard_names), set(loaded_metadata.shards))
         Assert.eq(loaded_metadata.shards[: len(shard_names)], list(shard_names))
 
         # Using `log_fn=bool` sets the output to true if the error list is non-empty.
@@ -95,7 +95,13 @@ def load(self, config: CheckpointLoadConfig) -> dict[str, typing.Any] | None:
                     )
                     path = config.path / f"rank_{rank}.safetensors"
                     log_main_rank(f"Loading from {path}", log_fn=logger.info)
-                    # TODO: skip shards without overlap.
+
+                    # First do a dry run to check if there is any overlap.
+                    if not self._has_shard_overlaps(loaded_model):
+                        # No overlap found, skip this file.
+                        continue
+
+                    # TODO: Lazy loading?
                     with safetensors.safe_open(path, framework="pt", device=str(self._model.distributed.device)) as f:
                         # TODO: Use self_shard
                         if "state_shard" in f.keys():
@@ -111,22 +117,36 @@ def load(self, config: CheckpointLoadConfig) -> dict[str, typing.Any] | None:
                                 shard_name: f.get_tensor(f"{shard_name}_shard") for shard_name in shard_names
                             }
 
-                        for shard_name, loaded_shard in loaded_shards.items():
-                            loaded_model.get_shard_meta(shard_name).validate(loaded_shard)
-
-                        self_shards = {shard_name: self._model.get_shard(shard_name) for shard_name in shard_names}
-
-                        counter = torch.zeros(1, dtype=torch.int64, device=self._model.distributed.device)
-                        for _, loaded_fsdp, loaded_fsdp_shards in loaded_model.split_shards_by_fsdp(loaded_shards):
-                            for _, self_fsdp, self_fsdp_shards in self._model.split_shards_by_fsdp(self_shards):
-                                self_fsdp.copy_shard_overlaps(
-                                    loaded_fsdp,
-                                    self_fsdp_shards,
-                                    loaded_fsdp_shards,
-                                    counter,
-                                    self._model.distributed.device,
-                                )
-
-                        context.mark_as_loaded(counter.item())
+                    self._copy_shard_overlaps(loaded_model, loaded_shards, context)
 
         return loaded_metadata.metadata
+
+    def _has_shard_overlaps(self, loaded_model) -> bool:
+        for _, loaded_fsdp, _ in loaded_model.split_shards_by_fsdp({}):
+            for _, self_fsdp, _ in self._model.split_shards_by_fsdp({}):
+                counter = self_fsdp.copy_shard_overlaps(
+                    loaded_fsdp,
+                    None,
+                    None,
+                    self._model.distributed.device,
+                )
+                if counter:
+                    return True
+        return False
+
+    def _copy_shard_overlaps(self, loaded_model, loaded_shards, context):
+        for shard_name, loaded_shard in loaded_shards.items():
+            loaded_model.get_shard_meta(shard_name).validate(loaded_shard)
+
+        self_shards = {shard_name: self._model.get_shard(shard_name) for shard_name in loaded_shards}
+
+        for _, loaded_fsdp, loaded_fsdp_shards in loaded_model.split_shards_by_fsdp(loaded_shards):
+            for _, self_fsdp, self_fsdp_shards in self._model.split_shards_by_fsdp(self_shards):
+                counter = self_fsdp.copy_shard_overlaps(
+                    loaded_fsdp,
+                    self_fsdp_shards,
+                    loaded_fsdp_shards,
+                    self._model.distributed.device,
+                )
+                for parameter, count in counter.items():
+                    context.mark_as_loaded(count, parameter, True)

fast_llm/engine/checkpoint/safe_load.py

@@ -5,9 +5,9 @@
 from torch.distributed import all_reduce
 
 from fast_llm.core.distributed import add_ephemeral_timeout
+from fast_llm.engine.multi_stage.config import ShardName
 from fast_llm.engine.multi_stage.fast_llm_model import FastLLMModel
 from fast_llm.functional.triton.pointwise import triton_fill
-from fast_llm.utils import Assert
 
 logger = logging.getLogger(__name__)
 
@@ -48,14 +48,17 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         if not exc_type:
             self._validate()
 
-    def mark_as_loaded(self, count: int, parameter: tuple[str, str] | None = None) -> None:
+    def mark_as_loaded(self, count: int, parameter: tuple[str, str] | None = None, partial: bool = False) -> None:
         self._loaded += count
         if parameter is not None:
             parameter_name, shard_name = parameter
             if shard_name not in self._loaded_parameters:
                 self._loaded_parameters[shard_name] = {}
-            Assert.not_incl(parameter_name, self._loaded_parameters[shard_name])
-            self._loaded_parameters[shard_name][parameter_name] = count
+            if not partial and parameter_name in self._loaded_parameters[shard_name]:
+                raise ValueError(f"Duplicate loaded parameter ({parameter_name}, {shard_name})")
+            self._loaded_parameters[shard_name][parameter_name] = (
+                self._loaded_parameters[shard_name].get(parameter_name, 0) + count
+            )
 
     def _validate(self) -> None:
         errors = []
@@ -105,7 +108,7 @@ def _check_missing(self, errors: list[str]) -> None:
                                 f"{missing_for_param:,} values missing out of {parameter.numel():,} for parameter {parameter_name} in stage {stage.index}, shard {shard_name}"
                                 f" (locally {local_missing_for_param:,} out of {local_values.numel():,})"
                             )
-                    missing_for_pad = buffer[-fsdp._global_pad :].isnan().sum().item()
+                    missing_for_pad = buffer[-fsdp._global_pad :].isnan().sum().item() if fsdp._global_pad > 0 else 0
                     if missing_for_pad > 0:
                         global_total += missing_for_pad
                         local_missing_for_pad = (
@@ -127,52 +130,53 @@ def _check_missing(self, errors: list[str]) -> None:
                 )
 
     def _check_parameters(self, errors: list[str]) -> None:
-        loaded_shard_names = set(self._loaded_parameters)
-        shard_names = set(self._self_shards)
-        if loaded_shard_names != shard_names:
-            errors.append(f"Incorrect loaded shards: {loaded_shard_names}!={shard_names}")
-        for shard_name in shard_names & loaded_shard_names:
-            counter_per_parameter = {
-                parameter_name: self._loaded_parameters[shard_name].pop(parameter_name, None)
-                for parameter_name in self._model.parameter_names
-            }
-            for parameter_name, count in self._loaded_parameters[shard_name].items():
-                errors.append(f'Loaded unknown parameter "{parameter_name}" for shard "{shard_name}" (count={count})')
-            for parameter_name, counter in counter_per_parameter.items():
-                if self._model.is_parameter_on_device(parameter_name):
-                    if counter is None:
-                        errors.append(f'Missing parameter "{parameter_name}" for shard "{shard_name}"')
-                elif counter is not None and counter > 0:
-                    errors.append(f'Loaded off-device parameter : "{parameter_name}" for shard "{shard_name}"')
-            if self._distributed.world_group is not None:
-                counter_list = []
-                for parameter_name, counter in counter_per_parameter.items():
-                    parameter_stage = self._model.get_parameter_stage(parameter_name)
-                    parameter_meta = parameter_stage.get_parameter_meta(parameter_name)
-                    if (
-                        counter is None
-                        or (not parameter_meta.is_tensor_parallel and self._distributed.config.tensor_rank != 0)
-                        or parameter_stage.is_tied_weight_copy
-                    ):
-                        # Ignore the counter from missing or duplicate tensors.
-                        counter = 0
-                    counter_list.append(counter)
-
-                counter_tensor = torch.tensor(counter_list, dtype=torch.int64).to(self._distributed.device)
-
-                add_ephemeral_timeout(self._distributed.world_group, self._timeout)
-                all_reduce(counter_tensor, group=self._distributed.world_group)
-                counter_per_parameter = {
-                    parameter_name: counter
-                    for parameter_name, counter in zip(counter_per_parameter, counter_tensor.tolist())
-                }
-            for parameter_name, counter in counter_per_parameter.items():
-                parameter_size = (
-                    self._model.get_parameter_stage(parameter_name)
-                    .get_parameter_meta(parameter_name)
-                    .global_shape.numel()
+        if set(self._loaded_parameters) != set(self._self_shards):
+            errors.append(f"Incorrect loaded shards: {tuple(self._loaded_parameters)}!={tuple(self._self_shards)}")
+
+        counters = []
+        # Compare local counts against expected values.
+        for stage, fsdp, parameter_name, parameter_meta in self._model.stages_fsdp_parameters:
+            for shard_name in self._self_shards if fsdp.requires_grad else [ShardName.weights]:
+                counter = self._loaded_parameters[shard_name].pop(parameter_meta.tensor_name, 0)
+                local_size = (
+                    fsdp.get_parameter_size_in_shard(parameter_name, shard_name)
+                    if self._model.is_parameter_on_device(parameter_name)
+                    else 0
                 )
+                if counter != local_size:
+                    errors.append(
+                        f'Local counter mismatch for parameter "{parameter_name}"'
+                        f' and shard "{shard_name}": loaded {counter}, expected {local_size}'
+                    )
+                # Accumulate in a list for global counter check.
+                if (
+                    not parameter_meta.is_tensor_parallel and self._distributed.config.tensor_rank != 0
+                ) or stage.is_tied_weight_copy:
+                    # Ignore the counter from duplicate tensors.
+                    counter = 0
+                counters.append(counter)
+
+        # Check for unexpected parameters.
+        for shard_name, loaded in self._loaded_parameters.items():
+            for parameter_name, count in loaded.items():
+                errors.append(f'Loaded unknown parameter "{parameter_name}" for shard "{shard_name}" (count={count})')
+
+        # All-reduce to get global counts.
+        if self._distributed.world_group is not None:
+            counter_tensor = torch.tensor(counters, dtype=torch.int64).to(self._distributed.device)
+            # This may be the first distributed barrier after loading, so we need to wait for everyone to finish.
+            add_ephemeral_timeout(self._distributed.world_group, self._timeout)
+            all_reduce(counter_tensor, group=self._distributed.world_group)
+            counters = counter_tensor.tolist()
+
+        # Compare global counts against expected values.
+        for stage, fsdp, parameter_name, parameter_meta in self._model.stages_fsdp_parameters:
+            for shard_name in self._self_shards if fsdp.requires_grad else [ShardName.weights]:
+                counter = counters.pop(0)
+                parameter_size = parameter_meta.global_shape.numel()
                 if counter != parameter_size:
                     errors.append(
-                        f'Global counter mismatch for parameter "{parameter_name}" and shard "{shard_name}": {counter} != {parameter_size}'
+                        f'Global counter mismatch for parameter "{parameter_name}"'
+                        f' and shard "{shard_name}": loaded {counter}, expected {parameter_size}'
                     )
+        assert not counters

fast_llm/engine/config_utils/tensor_space.py

@@ -147,7 +147,10 @@ def add_tensor_dim(self, dim: TensorDim) -> None:
         else:
             if dim.parallel_dim is not None:
                 assert dim.parallel_dim.name in self._distributed_config.distributed_dims, dim.parallel_dim.name
-                Assert.eq(dim.parallel_dim, self._distributed_config.distributed_dims[dim.parallel_dim.name])
+                Assert.eq(
+                    dim.parallel_dim.__dict__,
+                    self._distributed_config.distributed_dims[dim.parallel_dim.name].__dict__,
+                )
             self._tensor_dims[dim.name] = dim
 
     def get_tensor_dim(self, name: str) -> TensorDim: