core.safe_globals

Module Contents

Classes

SafeUnpickler

Restricted unpickler for FP8 extra-state checkpoints. Only allows the narrow set of types that _encode_extra_state can produce: plain Python containers, numeric scalars, and the PyTorch tensor/storage primitives used by pickle.dumps(tensor). Any attempt to instantiate a class outside this allowlist raises pickle.UnpicklingError, preventing arbitrary code execution via a crafted checkpoint.

Functions

register_safe_globals	Register megatron-core safe classes with torch serialization.
safe_load_from_bytes	Safe version (weights_only=True) of torch.storage._load_from_bytes.

Data

SAFE_GLOBALS

API

core.safe_globals.SAFE_GLOBALS

None

core.safe_globals.register_safe_globals()

Register megatron-core safe classes with torch serialization.

core.safe_globals.safe_load_from_bytes(b)

Safe version (weights_only=True) of torch.storage._load_from_bytes.

class core.safe_globals.SafeUnpickler

Bases: pickle.Unpickler

Restricted unpickler for FP8 extra-state checkpoints. Only allows the narrow set of types that _encode_extra_state can produce: plain Python containers, numeric scalars, and the PyTorch tensor/storage primitives used by pickle.dumps(tensor). Any attempt to instantiate a class outside this allowlist raises pickle.UnpicklingError, preventing arbitrary code execution via a crafted checkpoint.

Initialization

Initialize self. See help(type(self)) for accurate signature.

_SAFE_CLASSES: frozenset

‘frozenset(…)’

find_class(module: str, name: str)