Overlay: drop a checkpoint's own table tensors when a sidecar is attached, including the monolithic weight/weight_scale form — lets FP8-table checkpoints (nvidia, official FP8) serve on one card
Browse files- worker_image_quant.py +9 -0
worker_image_quant.py
CHANGED
|
@@ -673,6 +673,15 @@ class PleOffloadRunner:
|
|
| 673 |
matched_checkpoint_tensors += 1
|
| 674 |
if disk_complete_tables:
|
| 675 |
table = _ple_disk_shard_of(mapped_name)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 676 |
if table is not None and table.startswith(disk_complete_tables):
|
| 677 |
continue
|
| 678 |
yield weight_name, tensor
|
|
|
|
| 673 |
matched_checkpoint_tensors += 1
|
| 674 |
if disk_complete_tables:
|
| 675 |
table = _ple_disk_shard_of(mapped_name)
|
| 676 |
+
if table is None and mapped_name.endswith(
|
| 677 |
+
(".weight", ".weight_scale", ".weight_scale_2", ".input_scale")
|
| 678 |
+
):
|
| 679 |
+
# Checkpoints that ship the table as one tensor plus a
|
| 680 |
+
# global scale (the FP8 revision and the NVFP4 builds
|
| 681 |
+
# derived from it) instead of shard_N.weight parts. The
|
| 682 |
+
# sidecar owns those rows, so drop them before the loader
|
| 683 |
+
# tries to write a 51.2B tensor into the stub.
|
| 684 |
+
table = mapped_name.rsplit(".", 1)[0]
|
| 685 |
if table is not None and table.startswith(disk_complete_tables):
|
| 686 |
continue
|
| 687 |
yield weight_name, tensor
|