(checkpoint_path: str, rank: int, world_size: int)
| 49 | |
| 50 | |
| 51 | def split_checkpoint_files(checkpoint_path: str, rank: int, world_size: int) -> list[str]: |
| 52 | checkpoint_files = [ |
| 53 | os.path.join(checkpoint_path, f) |
| 54 | for f in filter(lambda x: x.endswith(".safetensors"), os.listdir(checkpoint_path)) |
| 55 | ] |
| 56 | files_per_rank = (len(checkpoint_files) + world_size - 1) // world_size |
| 57 | return checkpoint_files[rank * files_per_rank : (rank + 1) * files_per_rank] |
| 58 | |
| 59 | |
| 60 | def split_tensors(checkpoint_path: str, rank: int, world_size: int) -> dict[str, torch.Tensor]: |