Core model loading

Safe checkpoint reads, state-dict discovery, structural identity, lazy disk maps, and model construction.

Last updated: Oct 7, 2026

On this page

A safe local state-dict example Choosing the loading level Checkpoint trust boundary Complete reference

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

function

`assign_state_dict_strict`

Source

def assign_state_dict_strict(module: Any,state_dict: Mapping[str, Any],label: str = 'checkpoint') -> Any

worldfoundry.corefrom worldfoundry.core import assign_state_dict_strict

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

moduleAny

state_dictMapping[str, Any]

labelstrdefault 'checkpoint'

Returns

Any

function

`build_rename_dict`

Source

def build_rename_dict(source_state_dict,target_state_dict,split_qkv = False)

worldfoundry.corefrom worldfoundry.core import build_rename_dict

Print parameter-key matches between two state dicts for conversion scripts. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction).

Parameters

source_state_dict

target_state_dict

split_qkvdefault False

function

`hash_model_file`

Source

def hash_model_file(path,with_shape = True)

worldfoundry.corefrom worldfoundry.core import hash_model_file

Return an MD5 digest of checkpoint key names loaded from *path*. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction).

Parameters

path

with_shapedefault True

function

`hash_state_dict_keys`

Source

def hash_state_dict_keys(state_dict,with_shape = True)

worldfoundry.corefrom worldfoundry.core import hash_state_dict_keys

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_shapedefault True

Include tensor dimensions as well as parameter names.

Returns

MD5 hexadecimal digest used for loader-registry matching.

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.

function

`load_model`

Source

def 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,strict = True,post_load_hook = None)

worldfoundry.corefrom worldfoundry.core import load_model

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_dict or DiskMap`.

configdefault None

Keyword arguments passed to `model_class`.

torch_dtypedefault torch.bfloat16

Final model dtype.

devicedefault 'cpu'

Final model device when fine-grained VRAM management is absent.

state_dict_converterdefault None

Optional key/shape converter applied before assignment.

use_disk_mapdefault False

Read parameter tensors lazily instead of loading a full state dict.

module_mapdefault None

Source-class to VRAM-wrapper mapping. Enabling it delegates placement to `enable_vram_management`.

vram_configdefault None

Offload/onload/preparing/computation placement dictionary.

vram_limitdefault None

Optional used-memory limit in GiB for wrappers.

state_dictdefault None

Already-loaded weights; takes precedence over `path`.

strictdefault True

Require an exact parameter match. When false, construct a materialized module, retain its initialization for missing keys, and call `initialize(missing_keys)` when the model provides it.

post_load_hookdefault None

Optional in-place mutation called after checkpoint assignment and before VRAM wrappers or placement are installed.

Returns

An eval-mode model with weights assigned.

Notes

Construction uses meta-device initialization where possible. DeepSpeed ZeRO-3 receives its specialized state-dict assignment path.

function

`load_model_loader_registry`

Source

def load_model_loader_registry(config_path: str | Path,model_classes: Mapping[str, type]) -> ModelLoaderRegistry

worldfoundry.corefrom worldfoundry.core import load_model_loader_registry

Overview

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 | Path

model_classesMapping[str, type]

Returns

ModelLoaderRegistry

function

`load_model_with_disk_offload`

Source

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.corefrom worldfoundry.core import load_model_with_disk_offload

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`.

configdefault None

Keyword arguments passed to `model_class`.

torch_dtypedefault torch.bfloat16

Computation dtype.

devicedefault 'cpu'

Preparing and computation device.

state_dict_converterdefault None

Optional checkpoint-key converter.

module_mapdefault None

Required source-class to wrapper-class mapping.

Returns

Eval-mode model with disk-backed VRAM wrappers installed.

function

`load_safetensors_into_model_streaming`

Source

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.corefrom worldfoundry.core import load_safetensors_into_model_streaming

Overview

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

modelAny

checkpoint_dirstr | os.PathLike[str]

strictbooldefault True

devicestr | torch.devicedefault 'cpu'

dtypetorch.dtype | Nonedefault None

assignbool | Nonedefault None

Returns

dict[str, int]

function

`load_sharded_safetensors_parallel_with_progress`

Source

def load_sharded_safetensors_parallel_with_progress(checkpoint_dir: str)

worldfoundry.corefrom worldfoundry.core import load_sharded_safetensors_parallel_with_progress

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.safetensors or a model.safetensors.index.json plus its shards. A sibling .zst file is decompressed through the system zstd` command.

Returns

Merged state dictionary for all parameters named by the index.

Raises

FileNotFoundError: Neither the index nor fallback model file exists.
RuntimeError: External zstd decompression fails.

function

`load_state_dict`

Source

def load_state_dict(file_path,torch_dtype = None,device = 'cpu',pin_memory = False,verbose = 0)

worldfoundry.corefrom worldfoundry.core import load_state_dict

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_dtypedefault None

Optional dtype conversion applied to tensor values.

devicedefault 'cpu'

Device used while deserializing weights; CPU is the safe default.

pin_memorydefault False

Pin CPU tensors after loading to accelerate a later GPU copy.

verbosedefault 0

Print start/finish messages when at least `1`.

Returns

Flat state-dict mapping parameter names to tensors or checkpoint values.

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.

function

`load_state_dict_from_safetensors_index`

Source

def load_state_dict_from_safetensors_index(file_path,torch_dtype = None,device = 'cpu')

worldfoundry.corefrom worldfoundry.core import load_state_dict_from_safetensors_index

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_dtypedefault None

Optional dtype conversion for each tensor.

devicedefault 'cpu'

Device used by the safetensors reader.

Returns

Merged state dict containing all referenced shards.

Raises

ValueError: The index has no mapping-valued `weight_map`.

function

`load_torch_checkpoint`

Source

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) -> Any

worldfoundry.corefrom worldfoundry.core import load_torch_checkpoint

Overview

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_locationAnydefault 'cpu'

Destination understood by `torch.load`.

weights_onlybool | Nonedefault True

Safe deserialization mode. `None` omits the argument for compatibility with older PyTorch releases.

allow_unsafe_pickle_fallbackbooldefault False

Retry with unrestricted pickle only after a weights-only unpickling failure. Never enable for untrusted files.

kwargsAny

Additional `torch.load` options.

Returns

AnyObject produced by `torch.load`.

Warnings

Setting `allow_unsafe_pickle_fallback=True` can execute code embedded in a malicious checkpoint.

function

`load_torch_state_dict`

Source

def load_torch_state_dict(checkpoint_path: str | os.PathLike[str],map_location: Any = 'cpu') -> Any

worldfoundry.corefrom worldfoundry.core import load_torch_state_dict

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_locationAnydefault 'cpu'

Destination understood by `torch.load`.

Returns

AnyDeserialized checkpoint object. The function does not unwrap outer `state_dict/module keys; use load_state_dict` for that policy.

class

`ModelConfig`

Source

class ModelConfig(path: Optional[Union[str, list[str]]] = None,model_id: Optional[str] = None,origin_file_pattern: Optional[Union[str, list[str]]] = None,download_source: Optional[str] = None,local_model_path: Optional[str] = None,skip_download: Optional[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: Optional[Dict[str, torch.Tensor]] = None)

worldfoundry.corefrom worldfoundry.core import ModelConfig

Resolved model source, download policy, and VRAM placement settings. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction).

Attributes

pathOptional[Union[str, list[str]]]default None

Existing checkpoint path or paths. When set, no model-hub lookup is needed.

model_idOptional[str]default None

Hugging Face or ModelScope repository identifier.

origin_file_patternOptional[Union[str, list[str]]]default None

Optional allow-pattern within the model repository.

download_sourceOptional[str]default None

`"huggingface" or "modelscope"; defaults from WORLDFOUNDRY_DOWNLOAD_SOURCE` and then to Hugging Face.

local_model_pathOptional[str]default None

Local repository cache root. Defaults through `WORLDFOUNDRY_MODEL_DIR`.

skip_downloadOptional[bool]default None

Reuse local files without contacting the model hub.

offload_deviceOptional[Union[str, torch.device]]default None

Device used while weights are inactive.

offload_dtypeOptional[torch.dtype]default None

Dtype used while weights are inactive.

onload_deviceOptional[Union[str, torch.device]]default None

First-stage prefetch device.

onload_dtypeOptional[torch.dtype]default None

First-stage prefetch dtype.

preparing_deviceOptional[Union[str, torch.device]]default None

Second-stage prefetch device.

preparing_dtypeOptional[torch.dtype]default None

Second-stage prefetch dtype.

computation_deviceOptional[Union[str, torch.device]]default None

Device used for forward execution.

computation_dtypeOptional[torch.dtype]default None

Dtype used for forward execution.

clear_parametersbooldefault False

Compatibility flag for loaders that release source tensors after assignment.

state_dictOptional[Dict[str, torch.Tensor]]default None

Optional already-loaded weights, bypassing file loading.

Methods

methcheck_input()Source

Require either a concrete path or a model-hub identifier.

methparse_original_file_pattern()Source

Normalize the repository allow-pattern to a glob-compatible value.

methparse_download_source()Source

Resolve the configured model hub, including the environment override.

methparse_skip_download() -> boolSource

Overview

Resolve offline behavior from the field or environment.

Source docstring

Resolve offline behavior from the field or environment.

Accepts the common boolean spellings; an unparseable value raises instead of silently falling through to "download anyway", which would hang offline clusters whose users set e.g. `=1` expecting a skip.

Returns

bool

methdownload()Source

Download missing repository files into `local_model_path`.

methrequire_downloading(use_usp: bool = False)Source

Return whether this rank should contact the configured model hub.

Parameters

use_uspbooldefault False

methreset_local_model_path()Source

Overview

Apply the canonical WorldFoundry model directory when none is set.

Source docstring

Apply the canonical WorldFoundry model directory when none is set.

An explicitly configured `local_model_path always wins; WORLDFOUNDRY_MODEL_DIR is consumed by local_model_root_path()` and only fills the default (explicit argument > environment > default).

methdownload_if_necessary(use_usp: bool = False)Source

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_uspbooldefault False

methvram_config()Source

Return placement fields in the mapping expected by VRAM wrappers.

class

`ModelLoaderRegistry`

Source

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.corefrom worldfoundry.core import ModelLoaderRegistry

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.

function

`search_for_embeddings`

Source

def search_for_embeddings(state_dict)

worldfoundry.corefrom worldfoundry.core import search_for_embeddings

Return all tensor leaves from a nested state dict. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction).

Parameters

state_dict

function

`search_for_files`

Source

def search_for_files(folder,extensions)

worldfoundry.corefrom worldfoundry.core import search_for_files

Recursively find files matching any suffix in `extensions`. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction).

Parameters

folder

extensions

function

`search_parameter`

Source

def search_parameter(param,state_dict,atol = 0.001)

worldfoundry.corefrom worldfoundry.core import search_parameter

Find the first state-dict key whose tensor numerically matches `param`. Belongs to Core model loading (checkpoints, state dicts, DiskMap, construction).

Parameters

param

state_dict

atoldefault 0.001

function

`split_state_dict_with_prefix`

Source

def split_state_dict_with_prefix(state_dict)

worldfoundry.corefrom worldfoundry.core import split_state_dict_with_prefix

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.

Returns

List of sub-dictionaries in deterministic prefix order. Keys without a dot form their own prefix group.

function

`validate_state_dict_compatibility`

Source

def validate_state_dict_compatibility(module: Any,state_dict: Mapping[str, Any],label: str = 'checkpoint') -> None

worldfoundry.corefrom worldfoundry.core import validate_state_dict_compatibility

Overview

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

moduleAny

state_dictMapping[str, Any]

labelstrdefault 'checkpoint'

Returns

None