|
| 1 | +"""KunLunXin plugin implementations for filesystem async checkpointing.""" |
| 2 | + |
| 3 | +import importlib.util |
| 4 | +import inspect |
| 5 | +import logging |
| 6 | +import os |
| 7 | +import warnings |
| 8 | +from time import time |
| 9 | +from typing import List, Tuple, Union |
| 10 | + |
| 11 | +import torch |
| 12 | +from torch.distributed.checkpoint.filesystem import DEFAULT_SUFFIX, _StoragePrefix, _write_item |
| 13 | +from torch.distributed.checkpoint.planner import SavePlan, SavePlanner, WriteItemType |
| 14 | +from megatron.core.dist_checkpointing.strategies.filesystem_async import ( |
| 15 | + FileSystemWriterAsync, |
| 16 | + WriteBucket, |
| 17 | + _process_memory, |
| 18 | + _split_by_separation_hint, |
| 19 | + _split_by_size_and_type, |
| 20 | + get_write_results_queue, |
| 21 | + logger, |
| 22 | +) |
| 23 | + |
| 24 | + |
| 25 | +def _bridge_available() -> bool: |
| 26 | + """Check if megatron.bridge is available.""" |
| 27 | + return importlib.util.find_spec("megatron.bridge") is not None |
| 28 | + |
| 29 | + |
| 30 | +def prepare_write_data(self, plan: SavePlan, planner: SavePlanner) -> None: |
| 31 | + """Prepare async checkpoint write buckets for KunLunXin bridge path. |
| 32 | +
|
| 33 | + Args: |
| 34 | + plan: Save plan from PyTorch distributed planner. |
| 35 | + planner: Save planner for resolving bytes and tensor data. |
| 36 | + """ |
| 37 | + if not _bridge_available(): |
| 38 | + return FileSystemWriterAsync.prepare_write_data.__wrapped__(self, plan, planner) |
| 39 | + |
| 40 | + storage_plan: _StoragePrefix = plan.storage_data |
| 41 | + start = time() |
| 42 | + logger.debug(f"thread_count: {self.thread_count}, time: {start}") |
| 43 | + if self.separation_hint: |
| 44 | + assert self.thread_count > 1, "thread_count must be at least 2 if separation_hint is provided" |
| 45 | + bins = self.thread_count // 2 if self.separation_hint is not None else self.thread_count |
| 46 | + item_buckets = _split_by_size_and_type(bins, plan.items) |
| 47 | + logger.debug(f"bucket_prep, time: {time() - start}") |
| 48 | + |
| 49 | + start = time() |
| 50 | + file_count = 0 |
| 51 | + |
| 52 | + def gen_file(prefix=""): |
| 53 | + """Generate a unique checkpoint file name with optional prefix.""" |
| 54 | + nonlocal file_count |
| 55 | + file_name = f"{prefix}{storage_plan.prefix}{file_count}{DEFAULT_SUFFIX}" |
| 56 | + file_count += 1 |
| 57 | + return file_name |
| 58 | + |
| 59 | + def _clone_or_dequantize_if_needed(ten: torch.Tensor): |
| 60 | + """Detach and dequantize GPU tensors if needed.""" |
| 61 | + ten = ten.detach() |
| 62 | + if ten.device.type != "cpu" and ten.device.type == "cuda" and "dequantize" in type(ten).__dict__: |
| 63 | + ten = ten.dequantize() |
| 64 | + return ten |
| 65 | + |
| 66 | + self.write_buckets = [] |
| 67 | + for group_name, group_buckets in _split_by_separation_hint( |
| 68 | + item_buckets, self.separation_hint |
| 69 | + ).items(): |
| 70 | + for bucket in group_buckets: |
| 71 | + bytes_data = [ |
| 72 | + (item, planner.resolve_data(item)) |
| 73 | + for item in bucket |
| 74 | + if item.type == WriteItemType.BYTE_IO |
| 75 | + ] |
| 76 | + tensor_data = [ |
| 77 | + (item, _clone_or_dequantize_if_needed(planner.resolve_data(item))) |
| 78 | + for item in bucket |
| 79 | + if item.type != WriteItemType.BYTE_IO |
| 80 | + ] |
| 81 | + if len(bytes_data) > 0 or len(tensor_data) > 0: |
| 82 | + file_name = gen_file(prefix=group_name) |
| 83 | + self.write_buckets.append( |
| 84 | + ( |
| 85 | + os.path.join(self.checkpoint_dir, file_name), |
| 86 | + file_name, |
| 87 | + (bytes_data, tensor_data), |
| 88 | + ) |
| 89 | + ) |
| 90 | + |
| 91 | + if len(self.write_buckets) > 0: |
| 92 | + assert len(self.write_buckets) <= self.thread_count, ( |
| 93 | + len(self.write_buckets), |
| 94 | + self.thread_count, |
| 95 | + ) |
| 96 | + self.results_queue = get_write_results_queue() |
| 97 | + else: |
| 98 | + self.results_queue = None |
| 99 | + end = time() |
| 100 | + logger.debug(f"D2H and push, time: {end - start}") |
| 101 | + |
| 102 | + |
| 103 | +def preload_tensors_kunlunxin( |
| 104 | + write_buckets: List[WriteBucket], non_blocking=True |
| 105 | +) -> List[WriteBucket]: |
| 106 | + """Preload tensors with blocking D2H copies for KunLunXin. |
| 107 | +
|
| 108 | + Args: |
| 109 | + write_buckets: List of write buckets to preload. |
| 110 | + non_blocking: Whether to use non-blocking D2H (forced to False). |
| 111 | +
|
| 112 | + Returns: |
| 113 | + List of write buckets with tensors moved to CPU. |
| 114 | + """ |
| 115 | + if _bridge_available(): |
| 116 | + return write_buckets |
| 117 | + |
| 118 | + if non_blocking: |
| 119 | + warnings.warn( |
| 120 | + "D2H might fail when `non_blocking` is True, force it to be False for KunLunXin." |
| 121 | + ) |
| 122 | + return FileSystemWriterAsync.preload_tensors.__wrapped__(write_buckets, non_blocking=False) |
| 123 | + |
| 124 | + |
| 125 | +def write_preloaded_data( |
| 126 | + transform_list, |
| 127 | + local_proc_idx: int, |
| 128 | + write_bucket: WriteBucket, |
| 129 | + results_queue, |
| 130 | + count_queue, |
| 131 | + use_fsync: bool, |
| 132 | + **kwargs, |
| 133 | +) -> Union[Tuple[int, Exception], None]: |
| 134 | + """Write preloaded checkpoint data for KunLunXin bridge path. |
| 135 | +
|
| 136 | + Args: |
| 137 | + transform_list: Storage writer transforms. |
| 138 | + local_proc_idx: Index of the worker performing writing. |
| 139 | + write_bucket: Data to write to storage. |
| 140 | + results_queue: Queue to return write results (may be None for main worker). |
| 141 | + count_queue: Queue to signal task completion (may be None). |
| 142 | + use_fsync: If True, call os.fsync at end of saving. |
| 143 | +
|
| 144 | + Returns: |
| 145 | + Tuple of (proc_idx, results) or (proc_idx, exception), or None. |
| 146 | + """ |
| 147 | + if not _bridge_available(): |
| 148 | + return FileSystemWriterAsync.write_preloaded_data.__wrapped__( |
| 149 | + transform_list, |
| 150 | + local_proc_idx, |
| 151 | + write_bucket, |
| 152 | + results_queue, |
| 153 | + count_queue, |
| 154 | + use_fsync, |
| 155 | + **kwargs, |
| 156 | + ) |
| 157 | + |
| 158 | + logger = logging.getLogger(__name__) |
| 159 | + logger.debug(f"{local_proc_idx} started") |
| 160 | + mem_before = _process_memory() |
| 161 | + use_msc = kwargs.get("use_msc", False) |
| 162 | + |
| 163 | + local_results = [] |
| 164 | + try: |
| 165 | + file_name, storage_key, (bytes_data, tensor_data) = write_bucket |
| 166 | + extra_kwargs = {} |
| 167 | + if "serialization_format" in inspect.signature(_write_item).parameters: |
| 168 | + from torch.distributed.checkpoint.filesystem import SerializationFormat |
| 169 | + |
| 170 | + extra_kwargs["serialization_format"] = SerializationFormat.TORCH_SAVE |
| 171 | + if use_msc: |
| 172 | + import multistorageclient as msc |
| 173 | + |
| 174 | + open_file = msc.open |
| 175 | + else: |
| 176 | + open_file = open |
| 177 | + with open_file(file_name, "wb") as stream: |
| 178 | + for write_item, data in bytes_data: |
| 179 | + local_results.append( |
| 180 | + _write_item( |
| 181 | + *transform_list, stream, data, write_item, storage_key, **extra_kwargs |
| 182 | + ) |
| 183 | + ) |
| 184 | + |
| 185 | + for write_item, tensor in tensor_data: |
| 186 | + assert tensor.is_cpu |
| 187 | + if tensor.dtype != torch.bfloat16: |
| 188 | + tensor = tensor.bfloat16() |
| 189 | + if tensor.untyped_storage().size() != tensor.numel() * tensor.itemsize: |
| 190 | + tensor = tensor.clone() |
| 191 | + local_results.append( |
| 192 | + _write_item( |
| 193 | + *transform_list, stream, tensor, write_item, storage_key, **extra_kwargs |
| 194 | + ) |
| 195 | + ) |
| 196 | + |
| 197 | + if use_fsync: |
| 198 | + if use_msc: |
| 199 | + stream.fsync() |
| 200 | + else: |
| 201 | + os.fsync(stream.fileno()) |
| 202 | + local_output = (local_proc_idx, local_results) |
| 203 | + except Exception as e: |
| 204 | + logger.debug(f"{local_proc_idx} failed") |
| 205 | + local_output = (local_proc_idx, e) |
| 206 | + if results_queue is not None: |
| 207 | + results_queue.put(local_output) |
| 208 | + if count_queue is not None: |
| 209 | + count_queue.get() |
| 210 | + count_queue.task_done() |
| 211 | + |
| 212 | + mem_after = _process_memory() |
| 213 | + logger.debug( |
| 214 | + f"{local_proc_idx} consumed: {mem_after - mem_before}," |
| 215 | + f" before: {mem_before}, after: {mem_after}" |
| 216 | + ) |
| 217 | + return local_output |
| 218 | + |
0 commit comments