Add PyTorch .bin file validation to CI weight validation (#17533)

This commit is contained in:
Alison Shao
2026-01-24 19:18:15 -08:00
committed by GitHub
parent b829b797ef
commit 9121f22656

View File

@@ -1219,6 +1219,38 @@ def _validate_safetensors_file(file_path: str) -> bool:
return False
def _validate_pytorch_bin_file(file_path: str) -> bool:
"""
Validate that a PyTorch .bin file is readable and not corrupted.
This catches corruption issues like truncated downloads or invalid archives
that would cause errors like:
"RuntimeError: PytorchStreamReader failed reading file data/X: invalid header
or archive is corrupted"
Args:
file_path: Path to the .bin file
Returns:
True if the file is valid, False if corrupted
"""
try:
import torch
# Use weights_only=True for security and to avoid executing arbitrary code
# mmap=False to fully read the file and catch all corruption
torch.load(file_path, map_location="cpu", weights_only=True, mmap=False)
return True
except Exception as e:
logger.warning(
"Corrupted PyTorch bin file detected: %s - %s: %s",
file_path,
type(e).__name__,
str(e),
)
return False
def _check_index_files_exist(snapshot_dir: str) -> Tuple[bool, Optional[str]]:
"""
Check if all files listed in safetensors index files actually exist on disk.
@@ -1366,11 +1398,15 @@ def _validate_sharded_model(
[],
)
# Validate safetensors files for corruption
# Validate weight files for corruption
if group_info["suffix"] == "safetensors":
for f in group_info["files"]:
if not _validate_safetensors_file(f):
corrupted_files.append(f)
elif group_info["suffix"] == "bin":
for f in group_info["files"]:
if not _validate_pytorch_bin_file(f):
corrupted_files.append(f)
# Check for required index file for safetensors shards
if group_info["suffix"] == "safetensors":
@@ -1561,7 +1597,7 @@ def ci_validate_and_cleanup_local_snapshot(
)
return False
# Also validate single (non-sharded) safetensors files
# Also validate single (non-sharded) weight files
for f in local_weight_files:
base_name = os.path.basename(f)
# Check if this is a single model file (not sharded)
@@ -1580,6 +1616,21 @@ def ci_validate_and_cleanup_local_snapshot(
# Selective cleanup for single file
_cleanup_corrupted_files_selective(model_name_or_path, [f])
return False
# Also validate single PyTorch .bin files
elif base_name in [
"pytorch_model.bin",
"model.bin",
"adapter_model.bin",
]:
if not _validate_pytorch_bin_file(f):
log_info_on_rank0(
logger,
f"Corrupted model file {base_name} for {model_name_or_path}. "
"Will selectively clean and re-download this file.",
)
# Selective cleanup for single file
_cleanup_corrupted_files_selective(model_name_or_path, [f])
return False
return True
@@ -1613,12 +1664,15 @@ def _validate_weights_after_download(
if not weight_files:
return True # No weight files to validate
# Validate safetensors files
# Validate weight files (safetensors and .bin)
corrupted_files = []
for f in weight_files:
if f.endswith(".safetensors") and os.path.exists(f):
if not _validate_safetensors_file(f):
corrupted_files.append(os.path.basename(f))
elif f.endswith(".bin") and os.path.exists(f):
if not _validate_pytorch_bin_file(f):
corrupted_files.append(os.path.basename(f))
if corrupted_files:
# Clean up corrupted files so next attempt re-downloads them
@@ -1777,9 +1831,20 @@ def ci_validate_and_clean_hf_cache(model_path: str) -> None:
if not _validate_safetensors_file(sf_file):
corrupted_files.append(sf_file)
# Also find and validate PyTorch .bin files
bin_files = glob_module.glob(os.path.join(snapshot_dir, "*.bin"))
for bin_file in bin_files:
# Skip broken symlinks (os.path.exists returns False for them)
if not os.path.exists(bin_file):
continue
if not _validate_pytorch_bin_file(bin_file):
corrupted_files.append(bin_file)
if corrupted_files:
logger.warning(
"HFRunner: Found %d corrupted safetensors file(s) for %s. "
"HFRunner: Found %d corrupted weight file(s) for %s. "
"Removing to force re-download.",
len(corrupted_files),
model_path,