shrinath-suresh/blogs-docs-splitted huggingface dataset:
entrypoint is a function. Workers specified as a binary
     entrypoint do not canonically have a return value and the
     "return_values" field is meaningless and may be empty.
If "is_failed()" returns "True" then the "failures" field contains
   the failure information, again, mapped by the GLOBAL rank of the
   worker that failed.
The keys in "return_values" and "failures" are mutually exclusive,
   that is, a worker's final state can only be one of: succeeded,
   failed. Workers intentionally terminated by the agent according to
   the agent's restart policy, are not represented in either
   "return_values" nor "failures".
Watchdog in the Agent
A named pipe based watchdog can be enabled in "LocalElasticAgent" if
an environment variable "TORCHELASTIC_ENABLE_FILE_TIMER" with value 1
has been defined in the "LocalElasticAgent" process. Optionally,
another environment variable "TORCHELASTIC_TIMER_FILE" can be set
Expiration Timers
Expiration timers are set up on the same process as the agent and used
from your script to deal with stuck workers. When you go into a code-
block that has the potential to get stuck you can acquire an
expiration timer, which instructs the timer server to kill the process
if it does not release the timer by the self-imposed expiration
deadline.
Usage:
import torchelastic.timer as timer
   import torchelastic.agent.server as agent
def main():
       start_method = "spawn"
       message_queue = mp.get_context(start_method).Queue()
       server = timer.LocalTimerServer(message, max_interval=0.01)
       server.start() # non-blocking
   spec = WorkerSpec(
               fn=trainer_func,
               args=(message_queue,),
               ...<OTHER_PARAMS...>)
   agent = agent.LocalElasticAgent(spec, start_method)
   agent.run()

def trainer_func(message_queue):
       timer.configure(timer.LocalTimerClient(message_queue))
with timer.expires(after=60): # 60 second expiry
           # do some work
In the example above if "trainer_func" takes more than 60 seconds to
complete, then the worker process is killed and the agent retries the
worker group.
Client Methods
torch.distributed.elastic.timer.configure(timer_client)
Configures a timer client. Must be called before using "expires".
torch.distributed.elastic.timer.expires(after, scope=None, client=None)
Acquires a countdown timer that expires in "after" seconds from
   now, unless the code-block that it wraps is finished within the
   timeframe. When the timer expires, this worker is eligible to be
   reaped. The exact meaning of "reaped" depends on the client
   implementation. In most cases, reaping means to terminate the
   worker process. Note that the worker is NOT guaranteed to be reaped
   at exactly "time.now() + after", but rather the worker is
   "eligible" for being reaped and the "TimerServer" that the client
talks to will ultimately make the decision when and how to reap the
   workers with expired timers.
Usage:
  torch.distributed.elastic.timer.configure(LocalTimerClient())
  with expires(after=10):
      torch.distributed.all_reduce(...)

Server/Client Implementations
Below are the timer server and client pairs that are provided by
torchelastic.
Note:
Timer server and clients always have to be implemented and used in
  pairs since there is a messaging protocol between the server and
  client.
Below is a pair of timer server and client that is implemented based
on a "multiprocess.Queue".
class torch.distributed.elastic.timer.LocalTimerServer(mp_queue, max_interval=60, daemon=True)
Server that works with "LocalTimerClient". Clients are expected to
   be subprocesses to the parent process that is running this server.
   Each host in the job is expected to start its own timer server
   locally and each server instance manages timers for local workers
(running on processes on the same host).
class torch.distributed.elastic.timer.LocalTimerClient(mp_queue)
Client side of "LocalTimerServer". This client is meant to be used
   on the same host that the "LocalTimerServer" is running on and uses
   pid to uniquely identify a worker. This is particularly useful in
   situations where one spawns a subprocess (trainer) per GPU on a
   host with multiple GPU devices.
Below is another pair of timer server and client that is implemented
based on a named pipe.
class torch.distributed.elastic.timer.FileTimerServer(file_path, max_interval=10, daemon=True, log_event=None)
Server that works with "FileTimerClient". Clients are expected to
   be running on the same host as the process that is running this
   server. Each host in the job is expected to start its own timer
   server locally and each server instance manages timers for local
   workers (running on processes on the same host).
Parameters:
Parameters:
      * file_path (str) -- str, the path of a FIFO special file
        to be created.
  * **max_interval** (*float*) -- float, max interval in seconds
    for each watchdog loop.

  * **daemon** (*bool*) -- bool, running the watchdog thread in
    daemon mode or not. A daemon thread will not block a process
    to stop.

  * **log_event** (*Callable**[**[**str**,
    **Optional**[**FileTimerRequest**]**]**, **None**]*) --
    Callable[[Dict[str, str]], None], an optional callback for
    logging the events in JSON format.

class torch.distributed.elastic.timer.FileTimerClient(file_path, signal=Signals.SIGKILL)
Client side of "FileTimerServer". This client is meant to be used
   on the same host that the "FileTimerServer" is running on and uses
   pid to uniquely identify a worker. This client uses a named_pipe to
   send timer requests to the "FileTimerServer". This client is a
   producer while the "FileTimerServer" is a consumer. Multiple
clients can work with the same "FileTimerServer".
Parameters:
      * file_path (str) -- str, the path of a FIFO special file.
        "FileTimerServer" must have created it by calling os.mkfifo().
  * **signal** -- signal, the signal to use to kill the process.
    Using a negative or zero signal will not kill the process.

Writing a custom timer server/client
To write your own timer server and client extend the
"torch.distributed.elastic.timer.TimerServer" for the server and
"torch.distributed.elastic.timer.TimerClient" for the client. The
"TimerRequest" object is used to pass messages between the server and
client.
class torch.distributed.elastic.timer.TimerRequest(worker_id, scope_id, expiration_time)
Data object representing a countdown timer acquisition and release
   that is used between the "TimerClient" and "TimerServer". A
   negative "expiration_time" should be interpreted as a "release"
   request.
Note:
request.
Note:
 the type of "worker_id" is implementation specific. It is
 whatever the TimerServer and TimerClient implementations have on
 to uniquely identify a worker.

class torch.distributed.elastic.timer.TimerServer(request_queue, max_interval, daemon=True)
Entity that monitors active timers and expires them in a timely
   fashion. This server is responsible for reaping workers that have
   expired timers.
abstract clear_timers(worker_ids)
  Clears all timers for the given "worker_ids".

abstract get_expired_timers(deadline)
  Returns all expired timers for each worker_id. An expired timer
  is a timer for which the expiration_time is less than or equal
  to the provided deadline.

  Return type:
     *Dict*[str, *List*[*TimerRequest*]]

abstract register_timers(timer_requests)
  Processes the incoming timer requests and registers them with
  the server. The timer request can either be a acquire-timer or
release-timer request. Timer requests with a negative
      expiration_time should be interpreted as a release-timer
      request.
class torch.distributed.elastic.timer.TimerClient
Client library to acquire and release countdown timers by
   communicating with the TimerServer.
abstract acquire(scope_id, expiration_time)
  Acquires a timer for the worker that holds this client object
  given the scope_id and expiration_time. Typically registers the
  timer with the TimerServer.

abstract release(scope_id)
  Releases the timer for the "scope_id" on the worker this client
  represents. After this method is called, the countdown timer on
  the scope is no longer in effect.
Train script
If your train script works with "torch.distributed.launch" it will
continue working with "torchrun" with these differences:


No need to manually pass "RANK", "WORLD_SIZE", "MASTER_ADDR", and
   "MASTER_PORT".


"rdzv_backend" and "rdzv_endpoint" can be provided. For most users
   this will be set to "c10d" (see rendezvous). The default
   "rdzv_backend" creates a non-elastic rendezvous where
   "rdzv_endpoint" holds the master address.


Make sure you have a "load_checkpoint(path)" and
   "save_checkpoint(path)" logic in your script. When any number of
   workers fail we restart all the workers with the same program
   arguments so you will lose progress up to the most recent
   checkpoint (see elastic launch).


"use_env" flag has been removed. If you were parsing local rank by
   parsing the "--local_rank" option, you need to get the local rank
   from the environment variable "LOCAL_RANK" (e.g.
   "int(os.environ["LOCAL_RANK"])").
"int(os.environ["LOCAL_RANK"])").
Below is an expository example of a training script that checkpoints
on each epoch, hence the worst-case progress lost on failure is one
full epoch worth of training.
def main():
        args = parse_args(sys.argv[1:])
        state = load_checkpoint(args.checkpoint_path)
        initialize(state)
    # torch.distributed.run ensures that this will work
    # by exporting all the env vars needed to initialize the process group
    torch.distributed.init_process_group(backend=args.backend)

    for i in range(state.epoch, state.total_num_epochs)
         for batch in iter(state.dataset)
             train(batch, state.model)

         state.epoch += 1
         save_checkpoint(state)

For concrete examples of torchelastic-compliant train scripts, visit
our examples page.
Automatic Mixed Precision package - torch.amp
"torch.amp" provides convenience methods for mixed precision, where
some operations use the "torch.float32" ("float") datatype and other
operations use lower precision floating point datatype
("lower_precision_fp"): "torch.float16" ("half") or "torch.bfloat16".
Some ops, like linear layers and convolutions, are much faster in
"lower_precision_fp". Other ops, like reductions, often require the
dynamic range of "float32".  Mixed precision tries to match each op to
its appropriate datatype.
Ordinarily, "automatic mixed precision training" with datatype of
"torch.float16" uses "torch.autocast" and "torch.cuda.amp.GradScaler"
together, as shown in the CUDA Automatic Mixed Precision examples and
CUDA Automatic Mixed Precision recipe. However, "torch.autocast" and
"torch.cuda.amp.GradScaler" are modular, and may be used separately if
desired. As shown in the CPU example section of "torch.autocast",
"automatic mixed precision training/inference" on CPU with datatype of
"torch.bfloat16" only uses "torch.autocast".
For CUDA and CPU, APIs are also provided separately:


"torch.autocast("cuda", args...)" is equivalent to
  "torch.cuda.amp.autocast(args...)".


"torch.autocast("cpu", args...)" is equivalent to
  "torch.cpu.amp.autocast(args...)". For CPU, only lower precision
  floating point datatype of "torch.bfloat16" is supported for now.


Autocasting


Gradient Scaling


Autocast Op Reference


Op Eligibility


CUDA Op-Specific Behavior


CUDA Ops that can autocast to "float16"


CUDA Ops that can autocast to "float32"


CUDA Ops that promote to the widest input type


Prefer "binary_cross_entropy_with_logits" over
  "binary_cross_entropy"




CPU Op-Specific Behavior


CPU Ops that can autocast to "bfloat16"


CPU Ops that can autocast to "float32"


CPU Ops that promote to the widest input type




Autocasting
Autocasting
class torch.autocast(device_type, dtype=None, enabled=True, cache_enabled=None)
Instances of "autocast" serve as context managers or decorators
   that allow regions of your script to run in mixed precision.
In these regions, ops run in an op-specific dtype chosen by
   autocast to improve performance while maintaining accuracy. See the
   Autocast Op Reference for details.
When entering an autocast-enabled region, Tensors may be any type.
   You should not call "half()" or "bfloat16()" on your model(s) or
   inputs when using autocasting.
"autocast" should wrap only the forward pass(es) of your network,
   including the loss computation(s).  Backward passes under autocast
   are not recommended. Backward ops run in the same type that
   autocast used for corresponding forward ops.
Example for CUDA Devices:
  # Creates model and optimizer in default precision
  model = Net().cuda()
  optimizer = optim.SGD(model.parameters(), ...)
for input, target in data:
          optimizer.zero_grad()
      # Enables autocasting for the forward pass (model + loss)
      with autocast():
          output = model(input)
          loss = loss_fn(output, target)

      # Exits the context manager before backward()
      loss.backward()
      optimizer.step()

See the CUDA Automatic Mixed Precision examples for usage (along
   with gradient scaling) in more complex scenarios (e.g., gradient
   penalty, multiple models/losses, custom autograd functions).
"autocast" can also be used as a decorator, e.g., on the "forward"
   method of your model:
  class AutocastModel(nn.Module):
      ...
      @autocast()
      def forward(self, input):
          ...

Floating-point Tensors produced in an autocast-enabled region may
   be "float16". After returning to an autocast-disabled region, using
   them with floating-point Tensors of different dtypes may cause type
mismatch errors.  If so, cast the Tensor(s) produced in the
   autocast region back to "float32" (or other dtype if desired). If a
   Tensor from the autocast region is already "float32", the cast is a
   no-op, and incurs no additional overhead. CUDA Example:
  # Creates some tensors in default dtype (here assumed to be float32)
  a_float32 = torch.rand((8, 8), device="cuda")
  b_float32 = torch.rand((8, 8), device="cuda")
  c_float32 = torch.rand((8, 8), device="cuda")
  d_float32 = torch.rand((8, 8), device="cuda")

  with autocast():
      # torch.mm is on autocast's list of ops that should run in float16.
      # Inputs are float32, but the op runs in float16 and produces float16 output.
      # No manual casts are required.
      e_float16 = torch.mm(a_float32, b_float32)
      # Also handles mixed input types
      f_float16 = torch.mm(d_float32, e_float16)

  # After exiting autocast, calls f_float16.float() to use with d_float32
g_float32 = torch.mm(d_float32, f_float16.float())
CPU Training Example:
  # Creates model and optimizer in default precision
  model = Net()
  optimizer = optim.SGD(model.parameters(), ...)

  for epoch in epochs:
      for input, target in data:
          optimizer.zero_grad()

          # Runs the forward pass with autocasting.
          with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
              output = model(input)
              loss = loss_fn(output, target)

          loss.backward()
          optimizer.step()

CPU Inference Example:
  # Creates model in default precision
  model = Net().eval()

  with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
      for input in data:
          # Runs the forward pass with autocasting.
          output = model(input)

CPU Inference Example with Jit Trace:
  class TestModel(nn.Module):
      def __init__(self, input_size, num_classes):
super(TestModel, self).init()
              self.fc1 = nn.Linear(input_size, num_classes)
          def forward(self, x):
              return self.fc1(x)
  input_size = 2
  num_classes = 2
  model = TestModel(input_size, num_classes).eval()

  # For now, we suggest to disable the Jit Autocast Pass,
  # As the issue: https://github.com/pytorch/pytorch/issues/75956
  torch._C._jit_set_autocast_mode(False)

  with torch.cpu.amp.autocast(cache_enabled=False):
      model = torch.jit.trace(model, torch.randn(1, input_size))
  model = torch.jit.freeze(model)
  # Models Run
  for _ in range(3):
      model(torch.randn(1, input_size))

Type mismatch errors in an autocast-enabled region are a bug; if
   this is what you observe, please file an issue.
"autocast(enabled=False)" subregions can be nested in autocast-
   enabled regions. Locally disabling autocast can be useful, for
example, if you want to force a subregion to run in a particular
   "dtype".  Disabling autocast gives you explicit control over the
   execution type.  In the subregion, inputs from the surrounding
   region should be cast to "dtype" before use:
  # Creates some tensors in default dtype (here assumed to be float32)
  a_float32 = torch.rand((8, 8), device="cuda")
  b_float32 = torch.rand((8, 8), device="cuda")
  c_float32 = torch.rand((8, 8), device="cuda")
  d_float32 = torch.rand((8, 8), device="cuda")

  with autocast():
      e_float16 = torch.mm(a_float32, b_float32)
      with autocast(enabled=False):
          # Calls e_float16.float() to ensure float32 execution
          # (necessary because e_float16 was created in an autocasted region)
          f_float32 = torch.mm(c_float32, e_float16.float())

      # No manual casts are required when re-entering the autocast-enabled region.
torch.mm again runs in float16 and produces float16 output, regardless of input types.
      g_float16 = torch.mm(d_float32, f_float32)

The autocast state is thread-local.  If you want it enabled in a
   new thread, the context manager or decorator must be invoked in
   that thread.  This affects "torch.nn.DataParallel" and
   "torch.nn.parallel.DistributedDataParallel" when used with more
   than one GPU per process (see Working with Multiple GPUs).
Parameters:
      * device_type (str, required) -- Whether to use 'cuda'
        or 'cpu' device
  * **enabled** (*bool**, **optional*) -- Whether autocasting
    should be enabled in the region. Default: "True"

  * **dtype** (*torch_dtype**, **optional*) -- Whether to use
    torch.float16 or torch.bfloat16.

  * **cache_enabled** (*bool**, **optional*) -- Whether the weight
    cache inside autocast should be enabled. Default: "True"
class torch.cuda.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=True)
See "torch.autocast". "torch.cuda.amp.autocast(args...)" is
   equivalent to "torch.autocast("cuda", args...)"
torch.cuda.amp.custom_fwd(fwd=None, *, cast_inputs=None)
Helper decorator for "forward" methods of custom autograd functions
   (subclasses of "torch.autograd.Function").  See the example page
   for more detail.
Parameters:
      cast_inputs ("torch.dtype" or None, optional, default=None)
      -- If not "None", when "forward" runs in an autocast-enabled
      region, casts incoming floating-point CUDA Tensors to the target
      dtype (non-floating-point Tensors are not affected), then
      executes "forward" with autocast disabled. If "None",
      "forward"'s internal ops execute with the current autocast
      state.
Note:
 If the decorated "forward" is called outside an autocast-enabled
 region, "custom_fwd" is a no-op and "cast_inputs" has no effect.
torch.cuda.amp.custom_bwd(bwd)
Helper decorator for backward methods of custom autograd functions
   (subclasses of "torch.autograd.Function"). Ensures that "backward"
   executes with the same autocast state as "forward". See the example
   page for more detail.
class torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16, cache_enabled=True)
See "torch.autocast". "torch.cpu.amp.autocast(args...)" is
   equivalent to "torch.autocast("cpu", args...)"
Gradient Scaling
If the forward pass for a particular op has "float16" inputs, the
backward pass for that op will produce "float16" gradients. Gradient
values with small magnitudes may not be representable in "float16".
These values will flush to zero ("underflow"), so the update for the
corresponding parameters will be lost.
To prevent underflow, "gradient scaling" multiplies the network's
loss(es) by a scale factor and invokes a backward pass on the scaled
loss(es).  Gradients flowing backward through the network are then
scaled by the same factor.  In other words, gradient values have a
larger magnitude, so they don't flush to zero.
Each parameter's gradient (".grad" attribute) should be unscaled
before the optimizer updates the parameters, so the scale factor does
not interfere with the learning rate.
class torch.cuda.amp.GradScaler(init_scale=65536.0, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True)
get_backoff_factor()
  Returns a Python float containing the scale backoff factor.

get_growth_factor()
  Returns a Python float containing the scale growth factor.

get_growth_interval()
  Returns a Python int containing the growth interval.

get_scale()
  Returns a Python float containing the current scale, or 1.0 if
  scaling is disabled.

  Warning:

    "get_scale()" incurs a CPU-GPU sync.

is_enabled()
  Returns a bool indicating whether this instance is enabled.

load_state_dict(state_dict)
load_state_dict(state_dict)
  Loads the scaler state.  If this instance is disabled,
  "load_state_dict()" is a no-op.

  Parameters:
     **state_dict** (*dict*) -- scaler state.  Should be an object
     returned from a call to "state_dict()".

scale(outputs)
  Multiplies ('scales') a tensor or list of tensors by the scale
  factor.

  Returns scaled outputs.  If this instance of "GradScaler" is not
  enabled, outputs are returned unmodified.

  Parameters:
     **outputs** (*Tensor** or **iterable of Tensors*) -- Outputs
     to scale.

set_backoff_factor(new_factor)
  Parameters:
     **new_scale** (*float*) -- Value to use as the new scale
     backoff factor.

set_growth_factor(new_factor)
  Parameters:
     **new_scale** (*float*) -- Value to use as the new scale
     growth factor.

set_growth_interval(new_interval)
  Parameters:
     **new_interval** (*int*) -- Value to use as the new growth
interval.
state_dict()
  Returns the state of the scaler as a "dict".  It contains five
  entries:

  * ""scale"" - a Python float containing the current scale

  * ""growth_factor"" - a Python float containing the current
    growth factor

  * ""backoff_factor"" - a Python float containing the current
    backoff factor

  * ""growth_interval"" - a Python int containing the current
    growth interval

  * ""_growth_tracker"" - a Python int containing the number of
    recent consecutive unskipped steps.

  If this instance is not enabled, returns an empty dict.

  Note:

    If you wish to checkpoint the scaler's state after a
    particular iteration, "state_dict()" should be called after
    "update()".

step(optimizer, args, *kwargs)
  "step()" carries out the following two operations:

  1. Internally invokes "unscale_(optimizer)" (unless "unscale_()"
     was explicitly called for "optimizer" earlier in the
iteration).  As part of the "unscale_()", gradients are
         checked for infs/NaNs.
  2. If no inf/NaN gradients are found, invokes "optimizer.step()"
     using the unscaled gradients.  Otherwise, "optimizer.step()"
     is skipped to avoid corrupting the params.

  "*args" and "**kwargs" are forwarded to "optimizer.step()".

  Returns the return value of "optimizer.step(*args, **kwargs)".

  Parameters:
     * **optimizer** (*torch.optim.Optimizer*) -- Optimizer that
       applies the gradients.

     * **args** -- Any arguments.

     * **kwargs** -- Any keyword arguments.

  Warning:

    Closure use is not currently supported.

unscale_(optimizer)
  Divides ("unscales") the optimizer's gradient tensors by the
  scale factor.

  "unscale_()" is optional, serving cases where you need to modify
  or inspect gradients between the backward pass(es) and "step()".
  If "unscale_()" is not called explicitly,  gradients will be