diff --git a/fp8_to_nvfp4_streaming.py b/fp8_to_nvfp4_streaming.py index 16a56c6..186fd5f 100644 --- a/fp8_to_nvfp4_streaming.py +++ b/fp8_to_nvfp4_streaming.py @@ -358,12 +358,16 @@ class ShardCache: def load_weight_and_scale(cache: ShardCache, weight_map, name): - """Load an FP8 weight with its weight_scale_inv companion (if any).""" + """Load an FP8 weight with its scale companion (if any).""" weight = cache.get(weight_map[name]).get_tensor(name) scale_name = name.replace(".weight", ".scale") scale = None if scale_name in weight_map: - scale = cache.get(weight_map[scale_name]).get_tensor(scale_name) + try: + scale = cache.get(weight_map[scale_name]).get_tensor(scale_name) + except Exception: + # Scale listed in index but not in shard (BF16 weights have no scale) + pass return weight, scale