fix: fix PR issues

specture724 · specture724 · commit 9ce3f31677d7 · 2025-12-08T10:29:26.000Z
diff --git a/checkpoint_engine/ps.py b/checkpoint_engine/ps.py
@@ -93,7 +93,7 @@ class ParameterMeta(BaseModel):
     name: str
     dtype: _TorchDtype
     shape: _TorchSize
-    manually_aligned: bool = True
+    aligned_size: int
 
 
 class BucketRange(NamedTuple):
@@ -142,11 +142,7 @@ def _align_size(dtype: torch.dtype, shape: torch.Size) -> int:
 def _to_named_tensor(metas: list[ParameterMeta], offset: int = 0) -> list[dict]:
     ret = []
     for meta in metas:
-        size = (
-            _align_size(meta.dtype, meta.shape)
-            if meta.manually_aligned
-            else meta.dtype.itemsize * meta.shape.numel()
-        )
+        size = meta.aligned_size
         ret.append(
             {
                 "name": meta.name,
@@ -428,6 +424,7 @@ class TPMeta(BaseModel):
                     name=parameter_name,
                     shape=meta["shape"],
                     dtype=meta["dtype"],
+                    aligned_size=_align_size(meta["dtype"], meta["shape"]),
                 )
             tp_meta = tp_metas[parameter_name]
             if tp_meta.concat_dim != -1:
@@ -437,7 +434,10 @@ class TPMeta(BaseModel):
             shape = list(parameter_metas[name].shape)
             shape[tp_meta.concat_dim] = shape[tp_meta.concat_dim] * tp_meta.size
             parameter_metas[name] = ParameterMeta(
-                name=name, shape=torch.Size(shape), dtype=parameter_metas[name].dtype
+                name=name,
+                shape=torch.Size(shape),
+                dtype=parameter_metas[name].dtype,
+                aligned_size=_align_size(parameter_metas[name].dtype, torch.Size(shape)),
             )
         weights_in_cpu = [parameters_with_tp[name][key] for key in sorted(parameters_with_tp[name])]
         # TODO: here concat is serial, which may be slow
@@ -455,20 +455,15 @@ class TPMeta(BaseModel):
     return parameters
 
 
-def _register_checkpoint(
-    *,
-    files: list[str],
-    named_tensors: dict[str, torch.Tensor],
-    rank: int | None = None,
-) -> list[MemoryBuffer]:
-    logger.info(
-        f"[rank{rank}] start to register checkpoint with {len(files)} files and {len(named_tensors)} named_tensors"
-    )
-    if not files and not named_tensors:
-        return []
-    memory_buffers: list[MemoryBuffer] = []
+def _inplace_pin_memory(files: list[str], rank: int | None = None) -> list[MemoryBuffer]:
+    def _parse_and_pin_from_safetensors(file_path: str) -> MemoryBuffer:
+        """
+        safetensors format see https://huggingface.co/docs/safetensors/en/index#format.
+        We load the safetensors file as bytes, then parse the header manually to get parameter metas.
+        The actual tensor data is in the remaining bytes and is naturally aligned.
+        We pin the remaining bytes as the buffer, making pinning faster.
+        """
 
-    def inplace_pin_memory(files: list[str]) -> list[MemoryBuffer]:
         def _pin(t: torch.Tensor):
             """
             Pin the memory of tensor in-place.
@@ -478,138 +473,142 @@ def _pin(t: torch.Tensor):
             r = cudart.cudaHostRegister(t.data_ptr(), t.numel() * t.element_size(), 0)
             assert r == 0, f"pin memory error, error code: {r}"
 
-        def _inplace_pin_memory(file_path: str) -> MemoryBuffer:
-            """
-            safetensors format see https://huggingface.co/docs/safetensors/en/index#format.
-            We load the safetensors file as bytes, then parse the header manually to get parameter metas.
-            The actual tensor data is in the remaining bytes and is naturally aligned.
-            We pin the remaining bytes as the buffer, making pinning faster.
-            """
-            # TODO: should only support /dev/shm? but we found files in disk also work?
-            size = os.stat(file_path).st_size
-            flag_size = 8
-            t = torch.from_file(file_path, True, size, dtype=torch.uint8)
-            assert t.nbytes > flag_size, (
-                f"tensor nbytes {t.nbytes} should be greater than flag_size {flag_size}"
-            )
-            os.remove(file_path)
-            start_pos = (
-                int.from_bytes(t[0:flag_size].numpy().tobytes(), byteorder="little", signed=False)
-                + flag_size
-            )
-            header_tensor = t[flag_size:start_pos]
-            header = json.loads(header_tensor.numpy().tobytes())
-            if "__metadata__" in header:
-                header.pop("__metadata__")
+        # TODO: should only support /dev/shm? but we found files in disk also work?
+        size = os.stat(file_path).st_size
+        flag_size = 8
+        t = torch.from_file(file_path, True, size, dtype=torch.uint8)
+        assert t.nbytes > flag_size, (
+            f"tensor nbytes {t.nbytes} should be greater than flag_size {flag_size}"
+        )
+        start_pos = (
+            int.from_bytes(t[0:flag_size].numpy().tobytes(), byteorder="little", signed=False)
+            + flag_size
+        )
+        header_tensor = t[flag_size:start_pos]
+        header = json.loads(header_tensor.numpy().tobytes())
+        if "__metadata__" in header:
+            header.pop("__metadata__")
 
-            metas: list[ParameterMeta] = []
-            offset = 0
-            try:
-                for name, meta in sorted(header.items(), key=lambda x: x[1]["data_offsets"]):
-                    start, end = meta["data_offsets"]
-                    # safetensors format ensures offsets are aligned
-                    assert offset == start, f"offset {offset} should be equal to start {start}"
-                    metas.append(
-                        ParameterMeta(
-                            name=name,
-                            dtype=_getdtype(meta["dtype"]),
-                            shape=torch.Size(meta["shape"]),
-                            manually_aligned=False,
-                        )
-                    )
-                    offset = end
-            except Exception as e:
-                logger.error(f"fail to parse safetensors header from {file_path}: {e}")
-                raise
-
-            buffer = t[start_pos:]
-            assert offset == buffer.nbytes, (
-                f"offset {offset} should be equal to buffer.nbytes {buffer.nbytes}"
-            )
-            _pin(buffer)
-            return MemoryBuffer(buffer=buffer, size=buffer.nbytes, metas=metas)
-
-        local_memory_buffers: list[MemoryBuffer] = []
-        lock = threading.Lock()
-        idx = 0
-        with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor:
-            futures = [executor.submit(_inplace_pin_memory, file) for file in files]
-            for future in concurrent.futures.as_completed(futures):
-                memory_buffer = future.result()
-                with lock:
-                    local_memory_buffers.append(memory_buffer)
-                    logger.info(
-                        f"[rank{rank}] register pin_memory for file in /dev/shm {idx + 1}/{len(files)} finished"
+        metas: list[ParameterMeta] = []
+        offset = 0
+        try:
+            for name, meta in sorted(header.items(), key=lambda x: x[1]["data_offsets"]):
+                start, end = meta["data_offsets"]
+                # safetensors format ensures offsets are aligned
+                assert offset == start, f"offset {offset} should be equal to start {start}"
+                metas.append(
+                    ParameterMeta(
+                        name=name,
+                        dtype=_getdtype(meta["dtype"]),
+                        shape=torch.Size(meta["shape"]),
+                        aligned_size=end - start,
                     )
-                    idx += 1
-        return local_memory_buffers
-
-    def normal_pin_memory(
-        files: list[str], named_tensors: dict[str, torch.Tensor]
-    ) -> list[MemoryBuffer]:
-        parameters = _load_checkpoint(files)
-        if named_tensors:
-            parameters.update(named_tensors)
-        bucket_size = max(4 << 30, max(_align_size(x.dtype, x.shape) for x in parameters.values()))
-
-        class MemoryBucket(BaseModel):
-            size: int
-            metas: list[ParameterMeta]
-
-        buckets: list[MemoryBucket] = []
-        buckets.append(MemoryBucket(size=0, metas=[]))
-        for name, tensor in sorted(parameters.items()):
-            size = _align_size(tensor.dtype, tensor.shape)
-            if buckets[-1].size + size > bucket_size:
-                assert buckets[-1], f"buckets[{len(buckets) - 1}] should not be empty"
-                buckets.append(MemoryBucket(size=0, metas=[]))
-            buckets[-1].metas.append(
-                ParameterMeta(name=name, shape=tensor.shape, dtype=tensor.dtype)
-            )
-            buckets[-1].size += size
+                )
+                offset = end
+        except Exception as e:
+            logger.error(f"fail to parse safetensors header from {file_path}: {e}")
+            raise
 
-        local_memory_buffers = [
-            MemoryBuffer(buffer=torch.empty(0), size=bucket.size, metas=bucket.metas)
-            for bucket in buckets
-        ]
+        buffer = t[start_pos:]
+        assert offset == buffer.nbytes, (
+            f"offset {offset} should be equal to buffer.nbytes {buffer.nbytes}"
+        )
+        # Remove the file after successfully loading. This will avoid doubling the memory usage.
+        # We assume files in /dev/shm/ are temporary files. So it's safe to remove them after loading.
+        os.remove(file_path)
+        _pin(buffer)
+        logger.info(
+            f"[rank{rank}] inplace pin memory for file {file_path} finished, size {buffer.nbytes / 1024 / 1024:.2f}MiB"
+        )
+        return MemoryBuffer(buffer=buffer, size=buffer.nbytes, metas=metas)
 
-        def register_pin_memory(idx: int, size: int) -> tuple[int, torch.Tensor]:
-            buffer = torch.empty(size, dtype=torch.uint8, pin_memory=True)
-            return idx, buffer
-
-        def register_tensor(buffer: torch.Tensor, offset: int, tensor: torch.Tensor):
-            buffer[offset : offset + tensor.nbytes] = tensor.view(-1).view(dtype=torch.uint8)
-
-        with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor:
-            futures = [
-                executor.submit(register_pin_memory, idx, bucket.size)
-                for idx, bucket in enumerate(buckets)
-            ]
-            new_futures = []
-            for future in concurrent.futures.as_completed(futures):
-                idx, buffer = future.result()
-                assert buffer.numel() == buckets[idx].size, (
-                    f"buffer numel {buffer.numel()} should be equal to bucket size {buckets[idx].size}"
-                )
-                local_memory_buffers[idx].buffer = buffer
-                logger.info(
-                    f"[rank{rank}] register pin_memory for bucket {idx + 1}/{len(buckets)} finished, "
-                    f"size {buffer.numel() / 1024 / 1024:.2f}MiB, start to copy tensors to buffer"
+    local_memory_buffers: list[MemoryBuffer] = []
+    with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor:
+        local_memory_buffers = list(executor.map(_parse_and_pin_from_safetensors, files))
+    return local_memory_buffers
+
+
+def _normal_pin_memory(
+    files: list[str],
+    named_tensors: dict[str, torch.Tensor],
+    rank: int | None = None,
+) -> list[MemoryBuffer]:
+    parameters = _load_checkpoint(files)
+    if named_tensors:
+        parameters.update(named_tensors)
+    bucket_size = max(4 << 30, max(_align_size(x.dtype, x.shape) for x in parameters.values()))
+
+    class MemoryBucket(BaseModel):
+        size: int
+        metas: list[ParameterMeta]
+
+    buckets: list[MemoryBucket] = []
+    buckets.append(MemoryBucket(size=0, metas=[]))
+    for name, tensor in sorted(parameters.items()):
+        size = _align_size(tensor.dtype, tensor.shape)
+        if buckets[-1].size + size > bucket_size:
+            assert buckets[-1], f"buckets[{len(buckets) - 1}] should not be empty"
+            buckets.append(MemoryBucket(size=0, metas=[]))
+        buckets[-1].metas.append(
+            ParameterMeta(name=name, shape=tensor.shape, dtype=tensor.dtype, aligned_size=size)
+        )
+        buckets[-1].size += size
+
+    local_memory_buffers = [
+        MemoryBuffer(buffer=torch.empty(0), size=bucket.size, metas=bucket.metas)
+        for bucket in buckets
+    ]
+
+    def register_pin_memory(idx: int, size: int) -> tuple[int, torch.Tensor]:
+        buffer = torch.empty(size, dtype=torch.uint8, pin_memory=True)
+        return idx, buffer
+
+    def register_tensor(buffer: torch.Tensor, offset: int, tensor: torch.Tensor):
+        buffer[offset : offset + tensor.nbytes] = tensor.view(-1).view(dtype=torch.uint8)
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor:
+        futures = [
+            executor.submit(register_pin_memory, idx, bucket.size)
+            for idx, bucket in enumerate(buckets)
+        ]
+        new_futures = []
+        for future in concurrent.futures.as_completed(futures):
+            idx, buffer = future.result()
+            assert buffer.numel() == buckets[idx].size, (
+                f"buffer numel {buffer.numel()} should be equal to bucket size {buckets[idx].size}"
+            )
+            local_memory_buffers[idx].buffer = buffer
+            logger.info(
+                f"[rank{rank}] register pin_memory for bucket {idx + 1}/{len(buckets)} finished, "
+                f"size {buffer.numel() / 1024 / 1024:.2f}MiB, start to copy tensors to buffer"
+            )
+            offset = 0
+            for meta in buckets[idx].metas:
+                name = meta.name
+                tensor = parameters[name]
+                size = _align_size(tensor.dtype, tensor.shape)
+                assert size == _align_size(meta.dtype, meta.shape), (
+                    f"tensor {name} size {size} should be equal to meta size {_align_size(meta.dtype, meta.shape)}"
                 )
-                offset = 0
-                for meta in buckets[idx].metas:
-                    name = meta.name
-                    tensor = parameters[name]
-                    size = _align_size(tensor.dtype, tensor.shape)
-                    assert size == _align_size(meta.dtype, meta.shape), (
-                        f"tensor {name} size {size} should be equal to meta size {_align_size(meta.dtype, meta.shape)}"
-                    )
-                    new_futures.append(executor.submit(register_tensor, buffer, offset, tensor))
-                    offset += size
-            for future in concurrent.futures.as_completed(new_futures):
-                future.result()
-        return local_memory_buffers
+                new_futures.append(executor.submit(register_tensor, buffer, offset, tensor))
+                offset += size
+        for future in concurrent.futures.as_completed(new_futures):
+            future.result()
+    return local_memory_buffers
+
 
+def _register_checkpoint(
+    *,
+    files: list[str],
+    named_tensors: dict[str, torch.Tensor],
+    rank: int | None = None,
+) -> list[MemoryBuffer]:
+    logger.info(
+        f"[rank{rank}] start to register checkpoint with {len(files)} files and {len(named_tensors)} named_tensors"
+    )
+    if not files and not named_tensors:
+        return []
+    memory_buffers: list[MemoryBuffer] = []
     files_to_inplace_pin = [
         file
         for file in files
@@ -618,11 +617,10 @@ def register_tensor(buffer: torch.Tensor, offset: int, tensor: torch.Tensor):
     files_to_normal_pin = [file for file in files if file not in files_to_inplace_pin]
     if files_to_normal_pin or named_tensors:
         memory_buffers.extend(
-            normal_pin_memory(files=files_to_normal_pin, named_tensors=named_tensors)
+            _normal_pin_memory(files=files_to_normal_pin, named_tensors=named_tensors, rank=rank)
         )
     if files_to_inplace_pin:
-        memory_buffers.extend(inplace_pin_memory(files_to_inplace_pin))
-
+        memory_buffers.extend(_inplace_pin_memory(files_to_inplace_pin, rank=rank))
     return memory_buffers
 
 
@@ -671,11 +669,7 @@ def _gen_h2d_buckets(
         for idx, metas in enumerate(items.memory_buffer_metas_list):
             start_offset, offset = 0, 0
             for meta in metas.metas:
-                s = (
-                    _align_size(meta.dtype, meta.shape)
-                    if meta.manually_aligned
-                    else meta.dtype.itemsize * meta.shape.numel()
-                )
+                s = meta.aligned_size
                 if buckets[-1][1].size + s > bucket_size:
                     if offset - start_offset > 0:
                         buckets[-1][1].ranges.append(
@@ -1159,12 +1153,7 @@ def _detect_bucket_size(self, *, disable_h2d_buffer: bool = False) -> tuple[int,
         for items in self._current_global_parameter_metas.values():
             for metas_list in items.memory_buffer_metas_list:
                 for meta in metas_list.metas:
-                    max_tensor_bytes = max(
-                        max_tensor_bytes,
-                        _align_size(meta.dtype, meta.shape)
-                        if meta.manually_aligned
-                        else meta.dtype.itemsize * meta.shape.numel(),
-                    )
+                    max_tensor_bytes = max(max_tensor_bytes, meta.aligned_size)
         free_bytes_divided_3 = free_bytes // (3 * _ALIGN_SIZE) * _ALIGN_SIZE
         if max_tensor_bytes <= free_bytes_divided_3 and not disable_h2d_buffer:
             self._logger_rank0(f"[rank{self._rank}] use h2d buffer")
