chore: import upstream snapshot with attribution

2026-07-13 13:18:33 +08:00
commit 4ececc111a
2017 changed files with 331736 additions and 0 deletions
@@ -0,0 +1,287 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from deepspeed import comm as dist
+
+
+def print_rank_0(message):
+    if dist.get_rank() == 0:
+        print(message)
+
+
+class ContiguousMemoryAllocator(object):
+
+    def __init__(self, size, dtype, device):
+        self.buffer = torch.zeros(size, dtype=dtype, device=device)
+
+        #address to contiguous size available
+        self.contiguous_sizes = {}
+
+        self.contiguous_sizes[0] = size
+
+        #tensor id to its address
+        self.tensor_addresses = {}
+
+        #tensor address to its size
+        self.tensor_sizes = {}
+
+        #tensor address to ids
+        self.tensor_ids = {}
+
+        #id to tensors
+        self.tensor_map = {}
+
+        #id to params. Maps each tensor buffer to list of parameters that uses it
+        self.id_to_params = {}
+
+        self.total_size = size
+        self.total_free = size
+        self.largest_contiguous = size
+        self.max_allocated = 0
+
+        self.count = 0
+
+    #create a tensor of size from the pre-allocated buffer
+    #if not enough free space will fail
+    #if not enough contiguous space, will defragment and allocate
+    def allocate_tensor(self, size):
+        free_before = self.total_free
+
+        assert size <= self.total_free, "Not enough memory in buffer. Allocation failed"
+        if self.largest_contiguous < size:
+            print_rank_0("Needs defragmentation to allocate. Before Defragmentation:")
+            self.print_allocation(resolution=100)
+            self._defragment_memory()
+            #set the param data to the new tensor buffer locations
+            self._reset_param_data()
+            print_rank_0("After defragmentation:")
+            self.print_allocation(resolution=100)
+
+        self.total_free = self.total_free - size
+
+        allocated = self.total_size - self.total_free
+        if allocated > self.max_allocated:
+            self.max_allocated = allocated
+
+        tensor_address = self._get_new_tensor_address(size)
+
+        ret_tensor = self._get_new_tensor(tensor_address, size)
+        print_rank_0(
+            f"Free before allocation {free_before}. Allocating {size}. Free after allocation {self.total_free}. Max allocated {self.max_allocated}"
+        )
+        assert self.total_free + size == free_before, "Allocation bookkeeping error"
+
+        return ret_tensor
+
+    #assigns the tensor data to the param data and keeps track of the assignment
+    #any change the underlying buffer from defragmentation will cause a
+    #reassignment of the param data
+    def assign_to_param(self, tensor, param, numel, shape):
+        tensor_id = id(tensor)
+
+        assert tensor_id in self.tensor_map.keys(), "No such tensor allocated by the allocator."
+        assert tensor.numel() >= numel, "Assert tensor buffer does is not large enough"
+        assert tensor_id not in self.id_to_params.keys(), "This tensor has already been assigned to a param"
+
+        self.id_to_params[tensor_id] = [param]
+
+        replicated_tensor = tensor.narrow(0, 0, numel).view(shape)
+        param.data = replicated_tensor.data
+        param.contiguous_tensor_id = tensor_id
+
+    #deletes the tensor and frees up the underlying buffer
+    def release_tensor(self, tensor):
+        free_before = self.total_free
+        tensor_id = id(tensor)
+        tensor_size = tensor.numel()
+        self._release_tensor(tensor_id)
+        self._unassign_params(tensor_id)
+        self.total_free += tensor_size
+        print_rank_0(
+            f"Free before release {free_before}. Released {tensor.numel()}. Total free after {self.total_free}.")
+        assert self.total_free - tensor_size == free_before, "Release bookkeeping error"
+
+    def release_tensor_with_id(self, tensor_id):
+        free_before = self.total_free
+        assert tensor_id in self.tensor_map.keys(), "Invalid tensor id"
+        tensor = self.tensor_map[tensor_id]
+        tensor_size = tensor.numel()
+        self._release_tensor(tensor_id)
+        self._unassign_params(tensor_id)
+        self.total_free += tensor_size
+        print_rank_0(
+            f"Free before release {free_before}. Released {tensor.numel()}. Total free after {self.total_free}.")
+        assert self.total_free - tensor_size == free_before, "Release bookkeeping error"
+
+    #shows the current memory allocation at specified resolution
+    def print_allocation(self, resolution=200):
+        total_size = self.buffer.numel() * 1.0
+        empty = []
+        for addr, size in self.contiguous_sizes.items():
+            start = int(addr * resolution / total_size)
+            end = int((addr + size) * resolution / total_size)
+            empty.extend(range(start, end))
+        s = ''
+        for i in range(resolution):
+            s += '.' if i in empty else '|'
+        print_rank_0(s)
+
+    def max_allocated(self):
+        return self.max_allocated
+
+    #to be called after defragmentation that moves the tensor buffers
+    #this call reassigns the data of all the parameters using the tensor buffers
+    def _reset_param_data(self):
+        for id, tensor in self.tensor_map.items():
+            for param in self.id_to_params[id]:
+                param.data = tensor.narrow(0, 0, param.numel()).view(param.data.shape).data
+
+    def _unassign_params(self, tensor_id):
+        if tensor_id in self.id_to_params.keys():
+            del self.id_to_params[tensor_id]
+
+    def _release_tensor(self, tensor_id):
+        assert tensor_id in self.tensor_addresses, f"Tensor id {tensor_id} not found"
+
+        address = self.tensor_addresses[tensor_id]
+        contiguous_size = self.tensor_map[tensor_id].numel()
+
+        del self.tensor_addresses[tensor_id]
+        del self.tensor_ids[address]
+        del self.tensor_map[tensor_id]
+        del self.tensor_sizes[address]
+
+        self._consolidate_address(address, contiguous_size)
+        self.largest_contiguous = self._largest_contiguous()
+
+    def _consolidate_address(self, address, contiguous_size):
+
+        #consolidate next buffer
+        end_address = address + contiguous_size
+        if end_address in self.contiguous_sizes:
+            contiguous_size += self.contiguous_sizes[end_address]
+            del self.contiguous_sizes[end_address]
+
+        #consolidate previous buffer
+        for addr, size in self.contiguous_sizes.items():
+            if addr + size == address:
+                del self.contiguous_sizes[addr]
+                contiguous_size += size
+                address = addr
+                break
+
+        self.contiguous_sizes[address] = contiguous_size
+
+    def _defragment_memory(self):
+        empty_addresses = sorted(self.contiguous_sizes.keys())
+        tensor_addresses = sorted(self.tensor_addresses.values())
+
+        tensor_index = 0
+
+        while tensor_index < len(tensor_addresses):
+
+            empty_addr = empty_addresses[0]
+            empty_size = self.contiguous_sizes[empty_addr]
+
+            tensor_addr = tensor_addresses[tensor_index]
+            tensor_size = self.tensor_sizes[tensor_addr]
+            tensor_id = self.tensor_ids[tensor_addr]
+            tensor = self.tensor_map[self.tensor_ids[tensor_addr]]
+
+            assert tensor_size == tensor.numel(), \
+                f"Size mismatch. {tensor_size} is allocated at addr {tensor_addr} but tensor size is {tensor.numel()} "
+
+            assert empty_addr != tensor_addr, \
+                f"Cannot have same empty address {empty_addr} and tensor address {tensor_addr}"
+
+            if empty_addr < tensor_addr:
+
+                if empty_size >= tensor_size:
+                    dest_buffer = self.buffer.narrow(0, empty_addr, tensor_size)
+                    src_buffer = self.buffer.narrow(0, tensor_addr, tensor_size)
+                    dest_buffer.data.copy_(src_buffer.data)
+                else:
+
+                    #print_rank_0(f'empty addr : {empty_addr}, empty size {empty_size} tensor addr {tensor_addr} tensor size {tensor_size}')
+                    src_addr = tensor_addr
+                    dest_addr = empty_addr
+                    while src_addr < (tensor_addr + tensor_size):
+                        copy_size = min(empty_size, tensor_addr + tensor_size - src_addr)
+
+                        dest_buffer = self.buffer.narrow(0, dest_addr, copy_size)
+                        src_buffer = self.buffer.narrow(0, src_addr, copy_size)
+
+                        dest_buffer.data.copy_(src_buffer.data)
+
+                        src_addr += copy_size
+                        dest_addr += copy_size
+
+                self._replace_old_address_with_new(tensor_id, empty_addr)
+
+                tensor_index += 1
+
+            else:
+                tensor_index += 1
+
+            empty_addresses = sorted(self.contiguous_sizes.keys())
+
+    def _replace_old_address_with_new(self, tensor_id, new_address):
+
+        tensor = self.tensor_map[tensor_id]
+        tensor_size = tensor.numel()
+        tensor.data = self.buffer.narrow(0, new_address, tensor_size).data
+
+        self._release_tensor(tensor_id)
+        self._mark_as_occupied(new_address, tensor_size)
+
+        self.tensor_ids[new_address] = tensor_id
+        self.tensor_map[tensor_id] = tensor
+        self.tensor_addresses[tensor_id] = new_address
+        self.tensor_sizes[new_address] = tensor_size
+
+    def _get_new_tensor_address(self, size):
+        tensor_address = None
+        for address, contiguous_size in self.contiguous_sizes.items():
+            if contiguous_size >= size and \
+                    (tensor_address is None or \
+                    contiguous_size < self.contiguous_sizes[tensor_address]):
+                tensor_address = address
+        assert tensor_address is not None, "address cannot be None"
+        return tensor_address
+
+    def _get_new_tensor(self, address, size):
+        available_contiguous_size = self.contiguous_sizes[address]
+
+        assert size <= available_contiguous_size, \
+            f"Tensor numel {size} is large than available contiguous size {available_contiguous_size}"
+        self.count += 1
+        new_tensor = self.buffer.narrow(0, address, size)
+        tensor_id = id(new_tensor)
+        self.tensor_addresses[tensor_id] = address
+        self.tensor_sizes[address] = size
+
+        self.tensor_ids[address] = tensor_id
+        self.tensor_map[tensor_id] = new_tensor
+
+        self._mark_as_occupied(address, size)
+
+        return new_tensor
+
+    def _largest_contiguous(self):
+        if len(self.contiguous_sizes) > 0:
+            return max([size for _, size in self.contiguous_sizes.items()])
+        else:
+            return 0
+
+    def _mark_as_occupied(self, address, size):
+        available_contiguous_size = self.contiguous_sizes[address]
+        del self.contiguous_sizes[address]
+
+        if available_contiguous_size != size:
+            self.contiguous_sizes[address + size] = available_contiguous_size - size
+
+        self.largest_contiguous = self._largest_contiguous()