Return a mock object that satisfies dist.get_world_size / get_rank.

Collect gradients from model parameters using safe_get_full_grad API

Compile the module using the specified backend and kwargs. If a compiler_fn is set, it will be used instead of torch.compile().

Return torch.ops.deepspeed.<name>, raising clearly if not registered.

Compare gradients between DDP and DeepSpeed. Uses PyTorch's default tolerances for the tensor dtype (e.g., for bfloat16: rtol=1.6e-2, atol=1e

Returns the pre-allocated, padded output Tensor.

Collect gradients from DDP model

Given a DSModule key, attempt to instantiate

Overrides the internal iteration index

Advances an iterator along logical dimensions of matrix in units of whole tiles

Validate config constraints. Raises ValueError on invalid config.

Build the immutable per-stage folding spec from public config sizes.

Returns a pointer

Run a small CPU/Gloo distributed test without requiring visible GPU devices.

Returns whether access is valid or not

Checkpoint a model or part of the model. This has been directly copied from torch.utils.checkpoint.

Assemble and return the fp32 gradient of a low-precision (e.g., fp16) parameter. The return data type is that used for gradient accum

< Efficiently disables all accesses guarded by mask

r"""Coalesce ZeRO 1/2/3 gradient reduction across multiple engine.backward() calls. One with-block == one optimizer step: every backward insid

Recursively get the attribute of a module. Args: model (`torch.nn.Module`) The model to get the attribute from. m