(
self,
modules: list[torch.nn.Module],
offload_device: torch.device,
onload_device: torch.device,
offload_leader: torch.nn.Module,
onload_leader: torch.nn.Module | None = None,
parameters: list[torch.nn.Parameter] | None = None,
buffers: list[torch.Tensor] | None = None,
non_blocking: bool = False,
stream: torch.cuda.Stream | torch.Stream | None = None,
record_stream: bool | None = False,
low_cpu_mem_usage: bool = False,
onload_self: bool = True,
offload_to_disk_path: str | None = None,
group_id: int | str | None = None,
)
| 114 | |
| 115 | class ModuleGroup: |
| 116 | def __init__( |
| 117 | self, |
| 118 | modules: list[torch.nn.Module], |
| 119 | offload_device: torch.device, |
| 120 | onload_device: torch.device, |
| 121 | offload_leader: torch.nn.Module, |
| 122 | onload_leader: torch.nn.Module | None = None, |
| 123 | parameters: list[torch.nn.Parameter] | None = None, |
| 124 | buffers: list[torch.Tensor] | None = None, |
| 125 | non_blocking: bool = False, |
| 126 | stream: torch.cuda.Stream | torch.Stream | None = None, |
| 127 | record_stream: bool | None = False, |
| 128 | low_cpu_mem_usage: bool = False, |
| 129 | onload_self: bool = True, |
| 130 | offload_to_disk_path: str | None = None, |
| 131 | group_id: int | str | None = None, |
| 132 | ) -> None: |
| 133 | self.modules = modules |
| 134 | self.offload_device = offload_device |
| 135 | self.onload_device = onload_device |
| 136 | self.offload_leader = offload_leader |
| 137 | self.onload_leader = onload_leader |
| 138 | self.parameters = parameters or [] |
| 139 | self.buffers = buffers or [] |
| 140 | self.non_blocking = non_blocking or stream is not None |
| 141 | self.stream = stream |
| 142 | self.record_stream = record_stream |
| 143 | self.onload_self = onload_self |
| 144 | self.low_cpu_mem_usage = low_cpu_mem_usage |
| 145 | |
| 146 | self.offload_to_disk_path = offload_to_disk_path |
| 147 | self._is_offloaded_to_disk = False |
| 148 | |
| 149 | if self.offload_to_disk_path is not None: |
| 150 | # Instead of `group_id or str(id(self))` we do this because `group_id` can be "" as well. |
| 151 | self.group_id = group_id if group_id is not None else str(id(self)) |
| 152 | short_hash = _compute_group_hash(self.group_id) |
| 153 | self.safetensors_file_path = os.path.join(self.offload_to_disk_path, f"group_{short_hash}.safetensors") |
| 154 | |
| 155 | all_tensors = [] |
| 156 | for module in self.modules: |
| 157 | all_tensors.extend(list(module.parameters())) |
| 158 | all_tensors.extend(list(module.buffers())) |
| 159 | all_tensors.extend(self.parameters) |
| 160 | all_tensors.extend(self.buffers) |
| 161 | all_tensors = list(dict.fromkeys(all_tensors)) # Remove duplicates |
| 162 | |
| 163 | self.tensor_to_key = {tensor: f"tensor_{i}" for i, tensor in enumerate(all_tensors)} |
| 164 | self.key_to_tensor = {v: k for k, v in self.tensor_to_key.items()} |
| 165 | self.cpu_param_dict = {} |
| 166 | else: |
| 167 | self.cpu_param_dict = self._init_cpu_param_dict() |
| 168 | |
| 169 | self._torch_accelerator_module = ( |
| 170 | getattr(torch, torch.accelerator.current_accelerator().type) |
| 171 | if hasattr(torch, "accelerator") |
| 172 | else torch.cuda |
| 173 | ) |
nothing calls this directly
no test coverage detected