Index

A

alloc_storage() (in module olmo_core.utils) apply() (olmo_core.distributed.fsdp.FSDP method)

apply_to_tensors() (in module olmo_core.utils) auto_wrap() (olmo_core.distributed.fsdp.FSDP class method)

C

Checkpointer (class in olmo_core.distributed.checkpoint) chunk_unsharded() (olmo_core.distributed.tensors.ShardedFlatTensor method) clear_directory() (in module olmo_core.io)

clip_grad_norm_() (olmo_core.distributed.fsdp.FSDP method) compute_local_shape_and_global_offset() (in module olmo_core.distributed.tensors.dtensor_utils) compute_overlap_with() (olmo_core.distributed.checkpoint.TensorShardSpec method)

D

default_thread_count() (in module olmo_core.utils) deserialize_from_tensor() (in module olmo_core.io)

dir_is_empty() (in module olmo_core.io) dtype (olmo_core.distributed.checkpoint.TensorStorageMetadata attribute)

F

file_exists() (in module olmo_core.io) file_size() (in module olmo_core.io) flattened_offsets (olmo_core.distributed.checkpoint.TensorShardSpec attribute) forward() (olmo_core.distributed.fsdp.FSDP method)

free_storage() (in module olmo_core.utils) FSDP (class in olmo_core.distributed.fsdp) FSDPPrecision (class in olmo_core.distributed.fsdp) FSDPShardingStrategy (class in olmo_core.distributed.fsdp) FULL_SHARD (olmo_core.distributed.fsdp.FSDPShardingStrategy attribute)

G

gather() (olmo_core.distributed.tensors.ShardedFlatTensor method) gc_cuda() (in module olmo_core.utils) get_bytes_range() (in module olmo_core.io) get_default_device() (in module olmo_core.utils) get_flattened_offsets() (olmo_core.distributed.checkpoint.TensorShardSpec method) get_flattened_offsets_in_file() (olmo_core.distributed.checkpoint.TensorStorageMetadata method)

get_grad_norm() (in module olmo_core.utils) get_local_shape_and_global_offset() (in module olmo_core.distributed.tensors.dtensor_utils) get_merged_flattened_offsets() (olmo_core.distributed.checkpoint.TensorShardSpec method) get_metadata() (olmo_core.distributed.checkpoint.Checkpointer method) get_numel_in_file() (olmo_core.distributed.checkpoint.TensorStorageMetadata method) global_offset (olmo_core.distributed.checkpoint.TensorShardSpec attribute)

H

HYBRID_SHARD (olmo_core.distributed.fsdp.FSDPShardingStrategy attribute)

I

is_sharded (olmo_core.distributed.checkpoint.TensorStorageMetadata attribute) (olmo_core.distributed.tensors.ShardedFlatTensor property)

is_url() (in module olmo_core.io)

L

load() (olmo_core.distributed.checkpoint.Checkpointer method) load_model_and_optim_state() (in module olmo_core.distributed.checkpoint)

load_state_dict() (olmo_core.distributed.fsdp.FSDP method) local_numel (olmo_core.distributed.checkpoint.TensorShardSpec property) local_shape (olmo_core.distributed.checkpoint.TensorShardSpec attribute)

M

mark_as_sharded() (olmo_core.distributed.tensors.ShardedFlatTensor method) materialize_empty() (olmo_core.distributed.checkpoint.TensorStorageMetadata method) METADATA_FILENAME (olmo_core.distributed.checkpoint.Checkpointer attribute) metadata_set (olmo_core.distributed.tensors.ShardedFlatTensor property) module olmo_core.distributed.checkpoint olmo_core.distributed.fsdp olmo_core.distributed.tensors olmo_core.distributed.tensors.dtensor_utils olmo_core.exceptions olmo_core.io olmo_core.utils

module (olmo_core.distributed.fsdp.FSDP property)

N

named_buffers() (olmo_core.distributed.fsdp.FSDP method)

named_parameters() (olmo_core.distributed.fsdp.FSDP method)

O

olmo_core.distributed.checkpoint module olmo_core.distributed.fsdp module olmo_core.distributed.tensors module olmo_core.distributed.tensors.dtensor_utils module olmo_core.exceptions module

olmo_core.io module olmo_core.utils module OLMoEnvironmentError OLMoError OLMoNetworkError OLMoUserError

P

param_dtype (olmo_core.distributed.fsdp.FSDPPrecision attribute)

process_group (olmo_core.distributed.tensors.ShardedFlatTensor property)

R

reduce_dtype (olmo_core.distributed.fsdp.FSDPPrecision attribute)

reshard_() (olmo_core.distributed.tensors.ShardedFlatTensor method)

S

same_storage() (in module olmo_core.utils) save() (olmo_core.distributed.checkpoint.Checkpointer method) save_model_and_optim_state() (in module olmo_core.distributed.checkpoint) seed_all() (in module olmo_core.utils) serialize_to_tensor() (in module olmo_core.io) shape (olmo_core.distributed.checkpoint.TensorStorageMetadata attribute) shard() (olmo_core.distributed.tensors.ShardedFlatTensor class method) SHARD_GRAD_OP (olmo_core.distributed.fsdp.FSDPShardingStrategy attribute) shard_spec_per_file (olmo_core.distributed.checkpoint.TensorStorageMetadata attribute) sharded_chunk() (olmo_core.distributed.tensors.ShardedFlatTensor method) sharded_data (olmo_core.distributed.tensors.ShardedFlatTensor property)

sharded_numel (olmo_core.distributed.tensors.ShardedFlatTensor property) sharded_numels (olmo_core.distributed.tensors.ShardingSpec property) sharded_shape (olmo_core.distributed.tensors.ShardedFlatTensor property) ShardedFlatParameter (class in olmo_core.distributed.tensors) ShardedFlatTensor (class in olmo_core.distributed.tensors) sharding_spec (olmo_core.distributed.tensors.ShardedFlatTensor property) ShardingSpec (class in olmo_core.distributed.tensors) state_dict() (olmo_core.distributed.fsdp.FSDP method) StorageMetadata (class in olmo_core.distributed.checkpoint) StrEnum (class in olmo_core.utils) summon_full_params() (olmo_core.distributed.fsdp.FSDP method)

T

tensors (olmo_core.distributed.checkpoint.StorageMetadata attribute) TensorShardSpec (class in olmo_core.distributed.checkpoint)

TensorStorageMetadata (class in olmo_core.distributed.checkpoint) torch_dtype (olmo_core.distributed.checkpoint.TensorStorageMetadata property)

U

unshard() (olmo_core.distributed.checkpoint.Checkpointer method) unshard_() (olmo_core.distributed.tensors.ShardedFlatTensor method) unshard_model_state() (in module olmo_core.distributed.checkpoint) unshard_optim_state() (in module olmo_core.distributed.checkpoint) unsharded_data (olmo_core.distributed.tensors.ShardedFlatTensor property) unsharded_flattened_offsets (olmo_core.distributed.tensors.ShardedFlatTensor property) (olmo_core.distributed.tensors.ShardingSpec attribute)

unsharded_flattened_shape (olmo_core.distributed.tensors.ShardingSpec property) unsharded_numel (olmo_core.distributed.tensors.ShardedFlatTensor property) (olmo_core.distributed.tensors.ShardingSpec property) unsharded_shape (olmo_core.distributed.tensors.ShardedFlatTensor property) (olmo_core.distributed.tensors.ShardingSpec attribute) upload() (in module olmo_core.io)

W

wait_for() (in module olmo_core.utils)

wrap() (olmo_core.distributed.tensors.ShardedFlatTensor method) WRAPPED_MODULE_PREFIX (olmo_core.distributed.fsdp.FSDP attribute)