Contiguous memory buffer. Allocate a contiguous memory of type `dtype` and size `numel`. It is used to reduce memory fragmentation. Usage: After the allocation, the `_start` index is set tot the first index of the memory. A memory chunk starting from `_start` index
| 34 | |
| 35 | |
| 36 | class MemoryBuffer: |
| 37 | """Contiguous memory buffer. |
| 38 | Allocate a contiguous memory of type `dtype` and size `numel`. It is |
| 39 | used to reduce memory fragmentation. |
| 40 | |
| 41 | Usage: After the allocation, the `_start` index is set tot the first |
| 42 | index of the memory. A memory chunk starting from `_start` index |
| 43 | can be `allocated` for an input tensor, with the elements of the |
| 44 | tensor being coppied. The buffer can be reused by resetting the |
| 45 | `_start` index. |
| 46 | |
| 47 | """ |
| 48 | |
| 49 | def __init__(self, name, numel, dtype, track_usage): |
| 50 | if torch.distributed.get_rank() == 0: |
| 51 | element_size = torch.tensor([], dtype=dtype).element_size() |
| 52 | print( |
| 53 | "> building the {} memory buffer with {} num elements " |
| 54 | "and {} dtype ({:.1f} MB)...".format( |
| 55 | name, numel, dtype, numel * element_size / 1024 / 1024 |
| 56 | ), |
| 57 | flush=True, |
| 58 | ) |
| 59 | self.name = name |
| 60 | self.numel = numel |
| 61 | self.dtype = dtype |
| 62 | self.data = torch.empty( |
| 63 | self.numel, |
| 64 | dtype=self.dtype, |
| 65 | device=torch.cuda.current_device(), |
| 66 | requires_grad=False, |
| 67 | ) |
| 68 | |
| 69 | # Index tracking the start of the free memory. |
| 70 | self._start = 0 |
| 71 | |
| 72 | # Values used for tracking usage. |
| 73 | self.track_usage = track_usage |
| 74 | if self.track_usage: |
| 75 | self.in_use_value = 0.0 |
| 76 | self.total_value = 0.0 |
| 77 | |
| 78 | def reset(self): |
| 79 | """Reset the buffer start index to the beginning of the buffer.""" |
| 80 | self._start = 0 |
| 81 | |
| 82 | def is_in_use(self): |
| 83 | """Whether the current buffer hold on to any memory.""" |
| 84 | return self._start > 0 |
| 85 | |
| 86 | def numel_in_use(self): |
| 87 | """Return number of elements in use.""" |
| 88 | return self._start |
| 89 | |
| 90 | def add(self, tensor): |
| 91 | """Allocate a chunk of memory from the buffer to tensor and copy |
| 92 | the values.""" |
| 93 | assert ( |