Set the DEEPSPEED_AUTOTP_MODE based on the training flag

Store an iterator to sample for training data.

Sets a leaf flag for modules referenced by their names in ``model.named_modules()``. Args: model (torch.nn.Module): The model cont

Sets a leaf flag for modules referenced by suffixes of ``model.named_modules()`` names.

Args: max_log_level_str: maximum log level as a string Returns ``True`` if the current log_level is less or equal to the specified l

Return True when count_used_parameters_in_backward() should be re-evaluated.

Starts profiling. Extra attributes are added recursively to all the modules and the profiled torch.nn.functionals are monkey patched.

Update the model parameters. .. note:: This method will be called internally by ZeRO-Offload. DeepSpeed users should

Store a fragment to memory

Returns a torch.LongTensor that encodes partitioning information. Can be used along with ``data()`` to serialize a ``PartitionedTensor`` for

Compare two tensors or non-tensor numbers for their equality. Add msg=blah to add an additional comment to when assert fails.

r

Perform any necessary transformations of the gate parameter. Args: param (torch.Tensor): gate_w (shape: [num_experts, mo

Shards the first MLP parameter and passes it to the underlying implementation for any necessary transformations. Args:

Restore torch.autograd.Function and discard inputs captured for this compile cycle.

Return a preset carrying an actionable diagnostic for an unsupported HF model_type.

Update hook state after a gradient hook fires and run epilogue if all hooks have fired.

This method is identical to `logger.warning()`, but will emit the warning with the same message only once Note: The cache is for the functio

Returns True if the latest ``step()`` produced in parameter updates. Note that a ``False`` return is not an error condition. Steps are frequen

Zero FP16 parameter grads.

Pass through attributes defined in the model if they are not overridden by ds-engine.

A context manager to partition the model parameters during the model construction with MiCS partition strategy. Model states are partitioned

Create a mapping of n-dimensional tensor coordinates to linear indices. Arguments: axes (list): the names of the tensor axes