utils

class olmo_core.utils.StrEnum(value)

Bases: str, Enum

This is equivalent to Python’s enum.StrEnum since version 3.11. We include this here for compatibility with older version of Python.

olmo_core.utils.default_thread_count()

Return type:

int

olmo_core.utils.wait_for(condition, description, timeout=10.0)

Wait for the condition function to return True.

olmo_core.utils.apply_to_tensors(fn, container)

Recursively apply fn to all tensors in a container.

Return type:

None

olmo_core.utils.get_default_device()

Return type:

device

olmo_core.utils.seed_all(seed)

Seed all rng objects.

olmo_core.utils.get_grad_norm(params, norm_type)

Return the gradient norm of parameters, where the gradients are viewed as a single vector.

The returned norm is in FP32 even if parameters/gradients are in a low precision. This is because the downstream use of this return value is a reduction across ranks.

Return type:

Tensor

olmo_core.utils.same_storage(x, y)

Check if two tensors share the same storage.

Return type:

bool

olmo_core.utils.gc_cuda()

Run CUDA garbage collection.

olmo_core.utils.alloc_storage(tensor, size)

Allocate storage for tensor with the given size.

Returns True if this method allocated storage and False if the storage was already allocated.

Return type:

None

olmo_core.utils.free_storage(tensor)

Frees the underlying storage of tensor.

Returns True if the method freed the storage and False if the storage was already freed.

Return type:

None