(checkpoint_path, cpu_offload=False, remove_key=None, load_from_rank0=False, include_keys=None)
| 512 | |
| 513 | def load_weights(checkpoint_path, cpu_offload=False, remove_key=None, load_from_rank0=False, include_keys=None): |
| 514 | if not dist.is_initialized() or not load_from_rank0: |
| 515 | # Single GPU mode |
| 516 | logger.info(f"Loading weights from {checkpoint_path}") |
| 517 | cpu_weight_dict = load_pt_safetensors(checkpoint_path, remove_key, include_keys) |
| 518 | return cpu_weight_dict |
| 519 | |
| 520 | # Multi-GPU mode |
| 521 | is_weight_loader = False |
| 522 | current_rank = dist.get_rank() |
| 523 | if current_rank == 0: |
| 524 | is_weight_loader = True |
| 525 | |
| 526 | cpu_weight_dict = {} |
| 527 | if is_weight_loader: |
| 528 | logger.info(f"Loading weights from {checkpoint_path}") |
| 529 | cpu_weight_dict = load_pt_safetensors(checkpoint_path, remove_key) |
| 530 | |
| 531 | meta_dict = {} |
| 532 | if is_weight_loader: |
| 533 | for key, tensor in cpu_weight_dict.items(): |
| 534 | meta_dict[key] = {"shape": tensor.shape, "dtype": tensor.dtype} |
| 535 | |
| 536 | obj_list = [meta_dict] if is_weight_loader else [None] |
| 537 | |
| 538 | src_global_rank = 0 |
| 539 | dist.broadcast_object_list(obj_list, src=src_global_rank) |
| 540 | synced_meta_dict = obj_list[0] |
| 541 | |
| 542 | if cpu_offload: |
| 543 | target_device = "cpu" |
| 544 | distributed_weight_dict = {key: torch.empty(meta["shape"], dtype=meta["dtype"], device=target_device) for key, meta in synced_meta_dict.items()} |
| 545 | dist.barrier() |
| 546 | else: |
| 547 | target_device = torch.device(f"cuda:{current_rank}") |
| 548 | distributed_weight_dict = {key: torch.empty(meta["shape"], dtype=meta["dtype"], device=target_device) for key, meta in synced_meta_dict.items()} |
| 549 | dist.barrier(device_ids=[torch.cuda.current_device()]) |
| 550 | |
| 551 | for key in sorted(synced_meta_dict.keys()): |
| 552 | tensor_to_broadcast = distributed_weight_dict[key] |
| 553 | if is_weight_loader: |
| 554 | tensor_to_broadcast.copy_(cpu_weight_dict[key], non_blocking=True) |
| 555 | |
| 556 | if cpu_offload: |
| 557 | if is_weight_loader: |
| 558 | gpu_tensor = tensor_to_broadcast.cuda() |
| 559 | dist.broadcast(gpu_tensor, src=src_global_rank) |
| 560 | tensor_to_broadcast.copy_(gpu_tensor.cpu(), non_blocking=True) |
| 561 | del gpu_tensor |
| 562 | torch.cuda.empty_cache() |
| 563 | else: |
| 564 | gpu_tensor = torch.empty_like(tensor_to_broadcast, device="cuda") |
| 565 | dist.broadcast(gpu_tensor, src=src_global_rank) |
| 566 | tensor_to_broadcast.copy_(gpu_tensor.cpu(), non_blocking=True) |
| 567 | del gpu_tensor |
| 568 | torch.cuda.empty_cache() |
| 569 | else: |
| 570 | dist.broadcast(tensor_to_broadcast, src=src_global_rank) |
| 571 |
no test coverage detected