| 98 | |
| 99 | @dataclass |
| 100 | class GroupOffloadingConfig: |
| 101 | onload_device: torch.device |
| 102 | offload_device: torch.device |
| 103 | offload_type: GroupOffloadingType |
| 104 | non_blocking: bool |
| 105 | record_stream: bool |
| 106 | low_cpu_mem_usage: bool |
| 107 | num_blocks_per_group: int | None = None |
| 108 | offload_to_disk_path: str | None = None |
| 109 | stream: torch.cuda.Stream | torch.Stream | None = None |
| 110 | block_modules: list[str] | None = None |
| 111 | exclude_kwargs: list[str] | None = None |
| 112 | module_prefix: str = "" |
| 113 | |
| 114 | |
| 115 | class ModuleGroup: |
no outgoing calls
no test coverage detected
searching dependent graphs…