Buckets:
| """Invariants that must hold before committing a distributed SFT update.""" | |
| from __future__ import annotations | |
| import hashlib | |
| import torch | |
| import torch.distributed as dist | |
| def _buckets(parameters, limit_bytes=32 * 1024 * 1024): | |
| current, size = [], 0 | |
| for parameter in parameters: | |
| nbytes = parameter.numel() * parameter.element_size() | |
| if current and ( | |
| size + nbytes > limit_bytes | |
| or parameter.device != current[0].device | |
| or parameter.dtype != current[0].dtype | |
| ): | |
| yield current | |
| current, size = [], 0 | |
| current.append(parameter) | |
| size += nbytes | |
| if current: | |
| yield current | |
| def _collective(tensor, *, control_group, broadcast=False): | |
| # Unlike an asynchronous autograd hook, this boundary is reached only once | |
| # all local producers are finished. The CPU rendezvous keeps rank skew from | |
| # carrying the next bucket into an unfinished CUDA collective. | |
| if tensor.is_cuda: | |
| torch.cuda.synchronize(tensor.device) | |
| dist.barrier(group=control_group) | |
| if broadcast: | |
| dist.broadcast(tensor, src=0) | |
| else: | |
| dist.all_reduce(tensor, op=dist.ReduceOp.SUM) | |
| if tensor.is_cuda: | |
| torch.cuda.synchronize(tensor.device) | |
| def broadcast_trainable_parameters(parameters, *, control_group): | |
| """Initialize trainable replicas without registering DDP backward hooks.""" | |
| for bucket in _buckets(parameters): | |
| flat = torch.cat([p.detach().reshape(-1) for p in bucket]) | |
| _collective(flat, control_group=control_group, broadcast=True) | |
| offset = 0 | |
| for parameter in bucket: | |
| count = parameter.numel() | |
| parameter.copy_(flat[offset:offset + count].view_as(parameter)) | |
| offset += count | |
| def synchronize_gradients(parameters, *, control_group): | |
| """Sum already globally-normalized local gradients after accumulation. | |
| Every rank uses the same parameter schema and global used mask. Rank-local | |
| unused parameters contribute zero; globally unused parameters stay None so | |
| AdamW does not advance their momentum or apply weight decay. | |
| """ | |
| parameters = list(parameters) | |
| used = torch.tensor([p.grad is not None for p in parameters], dtype=torch.uint8) | |
| dist.all_reduce(used, op=dist.ReduceOp.MAX, group=control_group) | |
| active = [] | |
| for parameter, globally_used in zip(parameters, used.tolist(), strict=True): | |
| if not globally_used: | |
| parameter.grad = None | |
| continue | |
| if parameter.grad is None: | |
| parameter.grad = torch.zeros_like(parameter) | |
| if parameter.grad.is_sparse: | |
| raise ValueError("SFT post-backward synchronization requires dense gradients") | |
| active.append(parameter) | |
| for bucket in _buckets(active): | |
| flat = torch.cat([p.grad.reshape(-1) for p in bucket]) | |
| _collective(flat, control_group=control_group) | |
| offset = 0 | |
| for parameter in bucket: | |
| count = parameter.numel() | |
| parameter.grad.copy_(flat[offset:offset + count].view_as(parameter)) | |
| offset += count | |
| def assert_rank_equal(named_tensors, *, group, phase: str) -> None: | |
| """Reject unsynchronized replicas before an optimizer step or checkpoint. | |
| Exact equality is required: NCCL all-reduce distributes the same reduced | |
| buffer to all ranks. Reading detached CPU copies also fences device work. | |
| Only hashes travel over Gloo, not full model tensors. | |
| """ | |
| if group is None or dist.get_backend(group) != "gloo": | |
| raise ValueError("SFT replica checks require an explicit Gloo control group") | |
| digests = [] | |
| for name, tensor in named_tensors: | |
| if tensor is None: | |
| digests.append((name, None)) | |
| continue | |
| tensor = tensor.detach().cpu().contiguous() | |
| digest = hashlib.sha256() | |
| digest.update(str((tensor.shape, tensor.dtype)).encode()) | |
| digest.update(tensor.reshape(-1).view(torch.uint8).numpy().tobytes()) | |
| digests.append((name, digest.hexdigest())) | |
| gathered = [None] * dist.get_world_size(group) | |
| dist.all_gather_object(gathered, digests, group=group) | |
| if all(other == gathered[0] for other in gathered): | |
| return | |
| by_rank = [dict(values) for values in gathered] | |
| names = sorted(set().union(*(set(values) for values in by_rank))) | |
| differing = [ | |
| name | |
| for name in names | |
| if len({values.get(name, "absent") for values in by_rank}) > 1 | |
| ] | |
| raise RuntimeError( | |
| f"SFT rank consistency failed at {phase}: {len(differing)} tensors differ; " | |
| f"first tensors: {differing[:8]}. Refusing to commit unsynchronized training." | |
| ) | |
Xet Storage Details
- Size:
- 4.75 kB
- Xet hash:
- 1bdb4b034e564cb57b41dad3d424f79a7fb9aacf92e63fc15632fe01251c6299
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.