| 110 | backbone_weight_map: dict[str, str] = {} |
| 111 | |
| 112 | def flush_backbone_shard(): |
| 113 | nonlocal backbone_tensors, backbone_size, shard_idx |
| 114 | if not backbone_tensors: |
| 115 | return |
| 116 | shard_idx += 1 |
| 117 | shard_name = f"model-{shard_idx:05d}-of-PLACEHOLDER.safetensors" |
| 118 | shard_path = backbone_dir / shard_name |
| 119 | log.info(" Writing backbone shard %s (%d tensors, %.2f GB)", |
| 120 | shard_name, len(backbone_tensors), backbone_size / 1e9) |
| 121 | save_file(backbone_tensors, str(shard_path)) |
| 122 | for tname in backbone_tensors: |
| 123 | backbone_weight_map[tname] = shard_name |
| 124 | saved_shards.append(shard_name) |
| 125 | backbone_tensors = {} |
| 126 | backbone_size = 0 |
| 127 | |
| 128 | sorted_shards = sorted(shard_to_tensors.keys()) |
| 129 | for shard_file in sorted_shards: |