Core model loading
Safe checkpoint reads, state-dict discovery, structural identity, lazy disk maps, and model construction.
On this page
Model loading is layered so callers can choose how much policy Core owns. load_torch_checkpoint reads one PyTorch object with weights-only safety by default. load_state_dict understands files, folders, multiple paths, safetensors, and sharded indices. DiskMap exposes checkpoint tensors lazily by parameter name. load_model additionally constructs a module, converts keys, assigns weights, installs optional VRAM management, moves placement, and selects eval mode.
A safe local state-dict example
from pathlib import Path
from tempfile import TemporaryDirectory
import torch
from worldfoundry.core import hash_state_dict_keys, load_torch_state_dict
with TemporaryDirectory() as directory:
checkpoint = Path(directory) / "weights.pt"
expected = {"linear.weight": torch.arange(8).reshape(2, 4)}
torch.save(expected, checkpoint)
loaded = load_torch_state_dict(checkpoint, map_location="cpu")
assert torch.equal(loaded["linear.weight"], expected["linear.weight"])
print(hash_state_dict_keys(loaded, with_shape=True))hash_state_dict_keys is a routing fingerprint over key names and optionally shapes. It intentionally ignores values, so two checkpoints with the same architecture can share the digest. hash_model_file hashes bytes and is appropriate when content identity matters.
Choosing the loading level
Use load_torch_checkpoint when you know the file is a PyTorch checkpoint and need its original outer structure. Use load_state_dict when the input may be a directory, a sharded index, safetensors, or a list and you want one merged mapping. Use DiskMap when loading all tensors at once would exceed host memory; safetensors gives the best lazy behavior, while binary files use an in-memory compatibility reader.
Use load_model only when the shared construction path matches the model. It builds under a meta-device initialization context, supports a state-dict converter, handles the DeepSpeed ZeRO-3 assignment path, and can install module wrappers from module_map. A model with unusual parameter materialization should own that seam in its runner and use lower-level Core functions.
Checkpoint trust boundary
load_torch_checkpoint defaults to weights_only=True. The optional allow_unsafe_pickle_fallback=True can execute pickle payloads and must never be enabled for untrusted files. Safetensors does not carry this pickle execution risk. Remote URIs are localized through Core storage helpers before a reader is opened.
Complete reference
The blocks below are the generated signatures for this category. Use the on-page symbol index to jump; source links open the defining implementation behind each lazy export.
20 public symbols
def assign_state_dict_strict(module: Any,state_dict: Mapping[str, Any],label: str = 'checkpoint') -> Anyworldfoundry.core.assign_state_dict_strictfrom worldfoundry.core import assign_state_dict_strictOverview
Validate and assign tensors, including into a meta-device module. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction). Annotated return type: Any.
Parameters
moduleAnystate_dictMapping[str, Any]labelstr- default:
'checkpoint'
Returns: Any
def build_rename_dict(source_state_dict,target_state_dict,split_qkv = False)worldfoundry.core.build_rename_dictfrom worldfoundry.core import build_rename_dictOverview
Print parameter-key matches between two state dicts for conversion scripts. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction).
Parameters
source_state_dicttarget_state_dictsplit_qkv- default:
False
hash_model_file
funcdef hash_model_file(path, with_shape = True)worldfoundry.core.hash_model_filefrom worldfoundry.core import hash_model_fileOverview
Return an MD5 digest of checkpoint key names loaded from *path*. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction).
Parameters
pathwith_shape- default:
True
def hash_state_dict_keys(state_dict, with_shape = True)worldfoundry.core.hash_state_dict_keysfrom worldfoundry.core import hash_state_dict_keysOverview
Return a deterministic fingerprint of state-dict structure. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction).
Parameters
state_dict- Possibly nested mapping whose tensor keys identify a model checkpoint layout.
with_shape- Include tensor dimensions as well as parameter names.default:
True
Notes
This is an identity hint, not a content or security hash: tensor values are not read. Use `hash_model_file` when file content integrity is required.
load_model
funcdef load_model(model_class,path,config = None,torch_dtype = torch.bfloat16,device = 'cpu',state_dict_converter = None,use_disk_map = False,module_map = None,vram_config = None,vram_limit = None,state_dict = None)worldfoundry.core.load_modelfrom worldfoundry.core import load_modelOverview
Construct a model, assign checkpoint weights, and finalize inference placement. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction).
Parameters
model_class- PyTorch module class to instantiate.
path- Checkpoint path consumed by `
load_state_dictorDiskMap`. config- Keyword arguments passed to `
model_class`.default:None torch_dtype- Final model dtype.default:
torch.bfloat16 device- Final model device when fine-grained VRAM management is absent.default:
'cpu' state_dict_converter- Optional key/shape converter applied before assignment.default:
None use_disk_map- Read parameter tensors lazily instead of loading a full state dict.default:
False module_map- Source-class to VRAM-wrapper mapping. Enabling it delegates placement to `
enable_vram_management`.default:None vram_config- Offload/onload/preparing/computation placement dictionary.default:
None vram_limit- Optional used-memory limit in GiB for wrappers.default:
None state_dict- Already-loaded weights; takes precedence over `
path`.default:None
Notes
Construction uses meta-device initialization where possible. DeepSpeed ZeRO-3 receives its specialized state-dict assignment path.
def load_model_loader_registry(config_path: str | Path,model_classes: Mapping[str, type]) -> ModelLoaderRegistryworldfoundry.core.load_model_loader_registryfrom worldfoundry.core import load_model_loader_registryOverview
Load checkpoint-hash model routing rules from a package data YAML file. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction). Annotated return type: ModelLoaderRegistry.
Parameters
config_pathstr | Pathmodel_classesMapping[str, type]
Returns: ModelLoaderRegistry
def load_model_with_disk_offload(model_class,path,config = None,torch_dtype = torch.bfloat16,device = 'cpu',state_dict_converter = None,module_map = None)worldfoundry.core.load_model_with_disk_offloadfrom worldfoundry.core import load_model_with_disk_offloadOverview
Construct a model whose inactive weights remain disk-backed. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction).
Parameters
model_class- PyTorch module class to instantiate on the meta device.
path- Checkpoint path or paths indexed by `
DiskMap`. config- Keyword arguments passed to `
model_class`.default:None torch_dtype- Computation dtype.default:
torch.bfloat16 device- Preparing and computation device.default:
'cpu' state_dict_converter- Optional checkpoint-key converter.default:
None module_map- Required source-class to wrapper-class mapping.default:
None
def load_safetensors_into_model_streaming(model: Any,checkpoint_dir: str | os.PathLike[str],strict: bool = True,device: str | torch.device = 'cpu',dtype: torch.dtype | None = None,assign: bool | None = None) -> dict[str, int]worldfoundry.core.load_safetensors_into_model_streamingfrom worldfoundry.core import load_safetensors_into_model_streamingOverview
Load checkpoint shards one at a time and validate the complete key set. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction). Annotated return type: dict[str, int].
Source docstring
Load checkpoint shards one at a time and validate the complete key set.
Unlike a merged state-dict loader, peak host memory is bounded by the largest shard. Meta-initialized models are materialized tensor-by-tensor directly on `device; floating checkpoint tensors may also be converted to dtype` while streaming. Strict validation is performed after every shard has been applied, so the result is equivalent to one strict full-checkpoint load.
Parameters
modelAnycheckpoint_dirstr | os.PathLike[str]strictbool- default:
True devicestr | torch.device- default:
'cpu' dtypetorch.dtype | None- default:
None assignbool | None- default:
None
Returns: dict[str, int]
def load_sharded_safetensors_parallel_with_progress(checkpoint_dir: str)worldfoundry.core.load_sharded_safetensors_parallel_with_progressfrom worldfoundry.core import load_sharded_safetensors_parallel_with_progressOverview
Load a safetensors checkpoint, reading independent shards concurrently. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction).
Parameters
checkpoint_dirstr- Directory containing either `
model.safetensorsor amodel.safetensors.index.jsonplus its shards. A sibling.zstfile is decompressed through the systemzstd` command.
Raises
FileNotFoundError- Neither the index nor fallback model file exists.
RuntimeError- External zstd decompression fails.
load_state_dict
funcdef load_state_dict(file_path,torch_dtype = None,device = 'cpu',pin_memory = False,verbose = 0)worldfoundry.core.load_state_dictfrom worldfoundry.core import load_state_dictOverview
Load weights into a module with Core’s placement and key-handling policy.
Parameters
file_path- Local/remote checkpoint URI, directory, safetensors index, or list of any of those. Later files overwrite duplicate keys.
torch_dtype- Optional dtype conversion applied to tensor values.default:
None device- Device used while deserializing weights; CPU is the safe default.default:
'cpu' pin_memory- Pin CPU tensors after loading to accelerate a later GPU copy.default:
False verbose- Print start/finish messages when at least `
1`.default:0
Notes
File type is selected from the path: directories are scanned, `*.safetensors.index.json follows shards, *.safetensors` uses safetensors, and other suffixes use the safe PyTorch checkpoint loader.
def load_state_dict_from_safetensors_index(file_path,torch_dtype = None,device = 'cpu')worldfoundry.core.load_state_dict_from_safetensors_indexfrom worldfoundry.core import load_state_dict_from_safetensors_indexOverview
Load every unique shard referenced by a safetensors index. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction).
Parameters
file_path- Local or remote `
*.safetensors.index.json` URI. torch_dtype- Optional dtype conversion for each tensor.default:
None device- Device used by the safetensors reader.default:
'cpu'
Raises
ValueError- The index has no mapping-valued `
weight_map`.
def load_torch_checkpoint(checkpoint_path: str | os.PathLike[str],map_location: Any = 'cpu',weights_only: bool | None = True,allow_unsafe_pickle_fallback: bool = False,**kwargs: Any) -> Anyworldfoundry.core.load_torch_checkpointfrom worldfoundry.core import load_torch_checkpointOverview
Load a PyTorch checkpoint with weights-only deserialization by default. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction). Annotated return type: Any.
Parameters
checkpoint_pathstr | os.PathLike[str]- Local or remote checkpoint URI.
map_locationAny- Destination understood by `
torch.load`.default:'cpu' weights_onlybool | None- Safe deserialization mode. `
None` omits the argument for compatibility with older PyTorch releases.default:True allow_unsafe_pickle_fallbackbool- Retry with unrestricted pickle only after a weights-only unpickling failure. Never enable for untrusted files.default:
False kwargsAny- Additional `
torch.load` options.
Warnings
Setting `allow_unsafe_pickle_fallback=True` can execute code embedded in a malicious checkpoint.
Returns: Any — Object produced by `torch.load`.
def load_torch_state_dict(checkpoint_path: str | os.PathLike[str],map_location: Any = 'cpu') -> Anyworldfoundry.core.load_torch_state_dictfrom worldfoundry.core import load_torch_state_dictOverview
Load a PyTorch state-dict-shaped checkpoint in weights-only mode. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction). Annotated return type: Any.
Parameters
checkpoint_pathstr | os.PathLike[str]- Local or remote checkpoint URI.
map_locationAny- Destination understood by `
torch.load`.default:'cpu'
Returns: Any — Deserialized checkpoint object. The function does not unwrap outer `state_dict/module keys; use load_state_dict` for that policy.
ModelConfig
clsclass ModelConfig(path: Union[str, list[str]] = None,model_id: str = None,origin_file_pattern: Union[str, list[str]] = None,download_source: str = None,local_model_path: str = None,skip_download: bool = None,offload_device: Optional[Union[str, torch.device]] = None,offload_dtype: Optional[torch.dtype] = None,onload_device: Optional[Union[str, torch.device]] = None,onload_dtype: Optional[torch.dtype] = None,preparing_device: Optional[Union[str, torch.device]] = None,preparing_dtype: Optional[torch.dtype] = None,computation_device: Optional[Union[str, torch.device]] = None,computation_dtype: Optional[torch.dtype] = None,clear_parameters: bool = False,state_dict: Dict[str, torch.Tensor] = None)worldfoundry.core.ModelConfigfrom worldfoundry.core import ModelConfigOverview
Resolved model source, download policy, and VRAM placement settings. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction).
Attributes
pathUnion[str, list[str]]- Existing checkpoint path or paths. When set, no model-hub lookup is needed.default:
None model_idstr- Hugging Face or ModelScope repository identifier.default:
None origin_file_patternUnion[str, list[str]]- Optional allow-pattern within the model repository.default:
None download_sourcestr- `
"huggingface"or"modelscope"; defaults fromWORLDFOUNDRY_DOWNLOAD_SOURCE` and then to Hugging Face.default:None local_model_pathstr- Local repository cache root. Defaults through `
WORLDFOUNDRY_MODEL_DIR`.default:None skip_downloadbool- Reuse local files without contacting the model hub.default:
None offload_deviceOptional[Union[str, torch.device]]- Device used while weights are inactive.default:
None offload_dtypeOptional[torch.dtype]- Dtype used while weights are inactive.default:
None onload_deviceOptional[Union[str, torch.device]]- First-stage prefetch device.default:
None onload_dtypeOptional[torch.dtype]- First-stage prefetch dtype.default:
None preparing_deviceOptional[Union[str, torch.device]]- Second-stage prefetch device.default:
None preparing_dtypeOptional[torch.dtype]- Second-stage prefetch dtype.default:
None computation_deviceOptional[Union[str, torch.device]]- Device used for forward execution.default:
None computation_dtypeOptional[torch.dtype]- Dtype used for forward execution.default:
None clear_parametersbool- Compatibility flag for loaders that release source tensors after assignment.default:
False state_dictDict[str, torch.Tensor]- Optional already-loaded weights, bypassing file loading.default:
None
Methods
Overview
Require either a concrete path or a model-hub identifier.
Overview
Normalize the repository allow-pattern to a glob-compatible value.
Overview
Resolve the configured model hub, including the environment override.
Overview
Resolve offline behavior from the field or environment.
Overview
Download missing repository files into `local_model_path`.
Overview
Return whether this rank should contact the configured model hub.
Parameters
use_uspbool- default:
False
Overview
Apply the canonical WorldFoundry model directory when needed.
Overview
Materialize the configured source and replace `path` with local files.
Source docstring
Materialize the configured source and replace `path` with local files.
In USP execution only rank zero downloads; all ranks synchronize before the final local path is resolved.
Parameters
use_uspbool- default:
False
Overview
Return placement fields in the mapping expected by VRAM wrappers.
class ModelLoaderRegistry(model_loader_configs: list[tuple[Any, str, list[str], list[type], str]],huggingface_model_loader_configs: list[tuple[str, str, str, Any]],patch_model_loader_configs: list[tuple[str, list[str], list[type], dict[str, Any]]],preset_models_on_huggingface: dict[str, Any],preset_models_on_modelscope: dict[str, Any],preset_model_ids: tuple[str, ...],preset_model_websites: tuple[str, ...])worldfoundry.core.ModelLoaderRegistryfrom worldfoundry.core import ModelLoaderRegistryOverview
Validated routing data used to choose a model loader from checkpoint identity. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction).
Attributes
model_loader_configslist[tuple[Any, str, list[str], list[type], str]]- Single-file hash rules and their resolved model classes.
huggingface_model_loader_configslist[tuple[str, str, str, Any]]- Architecture-to-library routing rules.
patch_model_loader_configslist[tuple[str, list[str], list[type], dict[str, Any]]]- Patch/adaptor hash rules and extra kwargs.
preset_models_on_huggingfacedict[str, Any]- Named Hugging Face model presets.
preset_models_on_modelscopedict[str, Any]- Named ModelScope model presets.
preset_model_idstuple[str, ...]- Stable preset identifiers.
preset_model_websitestuple[str, ...]- Human-facing model source labels.
def search_for_embeddings(state_dict)worldfoundry.core.search_for_embeddingsfrom worldfoundry.core import search_for_embeddingsOverview
Return all tensor leaves from a nested state dict. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction).
Parameters
state_dict
search_for_files
funcdef search_for_files(folder, extensions)worldfoundry.core.search_for_filesfrom worldfoundry.core import search_for_filesOverview
Recursively find files matching any suffix in `extensions`. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction).
Parameters
folderextensions
search_parameter
funcdef search_parameter(param,state_dict,atol = 0.001)worldfoundry.core.search_parameterfrom worldfoundry.core import search_parameterOverview
Find the first state-dict key whose tensor numerically matches `param`. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction).
Parameters
paramstate_dictatol- default:
0.001
def split_state_dict_with_prefix(state_dict)worldfoundry.core.split_state_dict_with_prefixfrom worldfoundry.core import split_state_dict_with_prefixOverview
Split a state dict by the first dotted parameter-key segment. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction).
Parameters
state_dict- Mapping with string parameter names.
def validate_state_dict_compatibility(module: Any,state_dict: Mapping[str, Any],label: str = 'checkpoint') -> Noneworldfoundry.core.validate_state_dict_compatibilityfrom worldfoundry.core import validate_state_dict_compatibilityOverview
Reject missing, unexpected, or shape-incompatible checkpoint tensors. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction). Annotated return type: None.
Source docstring
Reject missing, unexpected, or shape-incompatible checkpoint tensors.
This validation is intentionally independent of tensor dtype: released FP32 weights may be assigned to a meta-device model and converted to the selected inference dtype during final placement.
Parameters
moduleAnystate_dictMapping[str, Any]labelstr- default:
'checkpoint'
Returns: None