(self, name, numel, dtype, track_usage)
| 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.""" |
nothing calls this directly
no outgoing calls
no test coverage detected