shrinath-suresh/blogs-docs-splitted huggingface dataset:
Perhaps the most important of all, in that it combines guards from
    all the provided "VariableTracker" instances passed in. It visits
    the guards and combines the guards from these onto itself.


It acts as a proxy on behalf of the underlying object, implementing
  methods for the rest of TorchDynamo to get information about the
  tracked object:


"call_method"


"call_function"


"python_type"


"as_proxy"


"is/as_python_proxy"


It stores the variable "source" of type "Source", from
  "torchdynamo/source.py". This source type is a relatively self
  contained class that helps us organize and bookkeep where the
  original source came from, and helps provide convenience methods for
  things like getting the name, and importantly for us, producing
  guards.


And this class ("VariableTracker") is built around subclassing,
somewhere between a full Abstract Base Class and fully fleshed out
class - it leaves many methods raising "NotImplementedError" - with
reliance on subclasses. See "torchdynamo/variables/" for all
subclasses to fulfill contracts and custom behaviors.
Knowing what we know now, we can see an example of how an instruction
from "dis", "BUILD_TUPLE":
"BUILD_TUPLE(count)" Creates a tuple consuming count items from the
   stack, and pushes the resulting tuple onto the stack.
In our case, our signature will be a little different due to the way
we create "Instruction" objects, but the gist of it will be the same.
Instead of passing in "count", we pass in an object with a little
extra bookkeeping, and of course, we deal with turning regular old
python objects into TorchDynamo notions:
def BUILD_TUPLE(self, inst):
       items = self.popn(inst.argval)
       options = VariableTracker.propagate(items)
       self.push(TupleVariable(items, **options))
Here is what this code does:

The function reads "argval", which in this case, is analogous to
   "counts" in the pydoc for the equivalent instruction.
The function "popn" the items, in this case, the signature is "def
   popn(self, n: int) -> List[TensorVariable]:" this hints at an
   underlying contract - we are returning "TensorVariables". If we
   take a closer look at "sybmolic_convert.py" and
   "InstructionTranslatorBase"/"InstructionTranslator"we see that the
   only thing pushed onto and popped from our stack are
   "VariableTracker"s.


The function calls "VariableTracker.propagate". This takes the
   guards from every single item popped off the stack in 2, and
   recursively traverses it and combines all the guards into
   "options": "py  return {      "guards": guards,  }"


The function then makes a new instance of a "VariableTracker",
   "TupleVariable"out of the "items" and "options". This then allows
   us to install all the appropriate guards from the "items" that make
   up the new "TupleVariable"


Note:
Where did the first guards come from? Propagation is a good
  technique, but we need something created before it can be
propagated. "VariableBuilder" calls "make_guards" as it creates
  "VariableTracker" instances, from "f_locals". This in turn calls
  into the "source", to have it create guards.
After all this, bytecode translation is done and we are one step
closer to producing "GuardedCode". We now understand how locals become
"VariableTracker"s, how instructions are handled, and where guards are
called on for creation. Before we can go into seeing how code and
guards are combined into a GuardedCode object, we need to dig a little
bit into those "make_guard" and "source.make_guard" calls above. We
can then understand, what was going on when we made guards alongside,
and on, "VariableTracker" instances.
Making Guards
Guards are just Python objects, of the class "Guard". Let's look at
them in more detail.
Looking at the definition of the dataclass (and therefore, ctor
signature), we see that it has a name, a source, and a create
function.
@dataclasses.dataclass
   class Guard:
       name: str
class Guard:
       name: str
       source: GuardSource
       create_fn: Callable
The name should be the name of the variable.
The source here is an enum indicating what kind of source the guard
belongs to.
Note:
Not to be confused with "Source" and the other types in "source.py",
  as stored on "VariableTracker".
"create_fn" provides the main functionality to transition from a
simple dataclass to actually producing valid Python code to be invoked
for knowing whether or not things have changed in between invocations,
and whether we can safely read from the code cache or not.
The most common code paths for getting an instance of a guard are
through "make_guards" on "VariableTracker".
"make_guards"->source.make_guard->return Guard(self.name(),
self.guard_source(), fn)
Or, in a concrete example:
...
   elif istype(value, range):
       guards = self.make_guards(GuardBuilder.EQUALS_MATCH)
       return RangeVariable(value=value, guards=guards)
Since "source" was set at the construction time of this
"VariableTracker", all that was needed here was to provide the "fn",
"GuardBuilder.EQUALS_MATCH" to the "create_fn" field.
This "create_fn" must be a method on "GuardBuilder". The reason for
this becomes apparent in our next step. Once we have all the guards
created for a frame, we move on to "CheckFunctionManager" and
"compile_check_fn".
Before the "convert_frame" function can produce a "GuardedCode", it
needs to run the "CheckFunctionManager", with all the guards, to
produce a "check_fn" which will then, in turn get passed in alongside
the code into "GuardedCode". This is the same "check_fn" that we store
in our cache entry, and the same one we run to know whether or not to
retrieve the code stored alongside. For reference, here is that code:
static CacheEntry create_cache_entry(CacheEntry next,
                                         PyObject guarded_code) {
     CacheEntry e = (CacheEntry *)malloc(sizeof(CacheEntry));
DEBUG_NULL_CHECK(e);
     e->check_fn = PyObject_GetAttrString(guarded_code, "check_fn");
     NULL_CHECK(e->check_fn);
     e->code = (PyCodeObject *)PyObject_GetAttrString(guarded_code, "code");
     NULL_CHECK(e->code);
     e->next = next;
     return e;
   }
We now know how a "check_fn" function is used, and who makes it, and
what it is composed of, but what we do not yet know is how. How does a
list of "Guard" objects become a function we can run later on?
First, we iterate these guards:
for guard in sorted(guards or [], key=Guard.sort_key):
       if not config.guard_nn_modules and guard.is_nn_module():
           continue
       guard.create(local_builder, global_builder)
Calling "guard.create" runs that "create_fn" we set on the "Guard"
class above (donât confuse it with the "check_fn" we are working on
producing, the names are similar, so it can get a little confusing).
In our example above, our "create_fn" is "GuardBuilder.EQUALS_MATCH".
So we are now invoking it, passing in the "self", the guard itself,
in.
The signature is: "def EQUALS_MATCH(self, guard: Guard):"
And internally to that function, we can use the "name" on the guard to
get back our original object, querying it for data and type
information, which in turn gets us to the most important bit:
appending code.
At its simplest, "EQUALS_MATCH" appends just one line of code:
"self.code.append(f"{ref} == {val!r}")". Where "ref" is the name of
the variable, and "val" is the value. It might produce code like this:
y == 2
This is a basic example. But if we append a few other kinds of
"GuardBuilder" functions and then combine them all with "and" in
between each statement (as we do), we might get something like this:
guardedcode.valid and _check_type_id(y, 94367738391392) and y == 2 and ___check_tensors(x)
Here is what this code performs:


A check for ".valid"


A type ID check


A value check


A tensor check
A value check


A tensor check


This becomes the heart of the code our "check_fn", which in turn is
evaluated the next time we encounter this code. It will then
check:


Is this code still valid?


If (1), Does "y" still have a type of "94367738391392"?


If (2), is "y" still 2?


If (3), letâs check on if tensor "x" changed in some specific ways.


If all of these are still true, then we can use the code cached
alongside this "check_fn".
Note:
For a deeper dive for how and where this happens you can read
  "static PyCodeObject lookup(CacheEntry e, PyObject *f_locals) {"
  of "_eval_frame.c".
If not, then, we can move on to recompiling the code anew, and storing
that in the cache alongside this code, and a whole new "check_fn",
again to be checked on yet another subsequent frame.
There are lots of other such functions on "GuardBuilder" which get
coalesced into, at times massive, strings which then get evaluated as
Python code and stored into "check_fn". The example above illustrates
of a simple case. To understand this functionality better, read the
other functions on "GuardBuilder", or better yet, dump the "code"
variable in "compile_check_fn" to see what is getting produced,
especially on larger, real models.
Summary
In this section, we have reviewed:


The role of ".valid" and invalidation around weak references (and
  potentially soon to be NN Moduleinvalidations).


How the C++ side of guard functions ("checktype_id",
  "_check_tensors", etc) operate


What happens when guards fail.


What happens if we produce invalid guard code.


We covered how user provided code wrapped in a TorchDynamo context
goes on to get traced and tracked internally, organized into
"VariableTracker"s "Source"s and subsequently "Guard"s, and how those
"Guards" in turn guide cache entry selection and invalidation when
handing Python code.
DDP Communication Hooks
DDP communication hook is a generic interface to control how to
communicate gradients across workers by overriding the vanilla
allreduce in DistributedDataParallel. A few built-in communication
hooks are provided, and users can easily apply any of these hooks to
optimize communication. Besides, the hook interface can also support
user-defined communication strategies for more advanced use cases.
How to Use a Communication Hook?
To use a communication hook, the user just needs to let the DDP model
register the hook before the training loop as below.
"torch.nn.parallel.DistributedDataParallel.register_comm_hook()"
What Does a Communication Hook Operate On?
A communication hook provides a flexible way to allreduce gradients.
Therefore, it mainly operates on the gradients on each replica before
allreduce, which are bucketized to increase the overlap between
communication and computation. Particularly,
"torch.distributed.GradBucket" represents a bucket of gradient tensors
to be allreduced.
class torch.distributed.GradBucket
This class mainly passes a flattened gradient tensor (returned by
   "buffer()") to DDP communication hook. This tensor can be further
   decomposed into a list of per-parameter tensors within this bucket
   (returned by "get_per_parameter_tensors()") to apply layer-wise
   operations.
torch.distributed.GradBucket.index(self: torch._C._distributed_c10d.GradBucket) -> int
Warning:
 Since the buckets are rebuilt after the first iteration, should
 not rely on the indices at the beginning of training.

Returns:
      The index of a bucket that stores gradients of a few contiguous
      layers. All the gradients are bucketized.
torch.distributed.GradBucket.buffer(self: torch._C._distributed_c10d.GradBucket) -> torch.Tensor
Returns:
      A flattened 1D "torch.Tensor" buffer, which can be further
decomposed into a list of per-parameter tensors within this
      bucket.
torch.distributed.GradBucket.gradients(self: torch._C._distributed_c10d.GradBucket) -> List[torch.Tensor]
Returns:
      A list of "torch.Tensor". Each tensor in the list corresponds to
      a gradient.
torch.distributed.GradBucket.is_last(self: torch._C._distributed_c10d.GradBucket) -> bool
Returns:
      Whether this bucket is the last bucket to allreduce in an
      iteration. This also means that this bucket corresponds to the
      first few layers in the forward pass.
torch.distributed.GradBucket.set_buffer(self: torch._C._distributed_c10d.GradBucket, buffer: torch.Tensor) -> None
Replaces the tensor in the bucket with the input tensor buffer.
torch.distributed.GradBucket.parameters(self: torch._C._distributed_c10d.GradBucket) -> List[torch.Tensor]
Returns:
      A list of "torch.Tensor". Each tensor in the list corresponds to
      a model parameter.
Default Communication Hooks
Default Communication Hooks
Default communication hooks are simple stateless hooks, so the
input state in "register_comm_hook" is either a process group or
"None". The input "bucket" is a "torch.distributed.GradBucket" object.
torch.distributed.algorithms.ddp_comm_hooks.default_hooks.allreduce_hook(process_group, bucket)
This DDP communication hook just calls "allreduce" using
   "GradBucket" tensors. Once gradient tensors are aggregated across
   all workers, its "then" callback takes the mean and returns the
   result. If user registers this hook, DDP results is expected to be
   same as the case where no hook was registered. Hence, this won't
   change behavior of DDP and user can use this as a reference or
   modify this hook to log useful information or any other purposes
   while unaffecting DDP behavior.
Example::
      >>> ddp_model.register_comm_hook(process_group, allreduce_hook)
Return type:
      Future[Tensor]
Return type:
      Future[Tensor]
torch.distributed.algorithms.ddp_comm_hooks.default_hooks.fp16_compress_hook(process_group, bucket)
This DDP communication hook implements a simple gradient
   compression approach that casts "GradBucket" tensor to half-
   precision floating-point format ("torch.float16") and then divides
   it by the process group size. It allreduces those "float16"
   gradient tensors. Once compressed gradient tensors are allreduced,
   the chained callback "decompress" casts it back to the input data
   type (such as "float32").
Example::
      >>> ddp_model.register_comm_hook(process_group, fp16_compress_hook)
Return type:
      Future[Tensor]
torch.distributed.algorithms.ddp_comm_hooks.default_hooks.bf16_compress_hook(process_group, bucket)
Warning: This API is experimental, and it requires NCCL version
   later than 2.9.6.
This DDP communication hook implements a simple gradient
   compression approach that casts "GradBucket" tensor to half-
precision Brain floating point format ("torch.bfloat16") and then
   divides it by the process group size. It allreduces those
   "bfloat16" gradient tensors. Once compressed gradient tensors are
   allreduced, the chained callback "decompress" casts it back to the
   input data type (such as "float32").
Example::
      >>> ddp_model.register_comm_hook(process_group, bf16_compress_hook)
Return type:
      Future[Tensor]
Additionally, a communication hook wrapper is provided to support
"fp16_compress_hook()" or "bf16_compress_hook()" as a wrapper, which
can be combined with other communication hooks.
torch.distributed.algorithms.ddp_comm_hooks.default_hooks.fp16_compress_wrapper(hook)
This wrapper casts the input gradient tensor of a given DDP
   communication hook to half-precision floating point format
   ("torch.float16"), and casts the resulting tensor of the given hook
   back to the input data type, such as "float32".
Therefore, "fp16_compress_hook" is equivalent to
"fp16_compress_wrapper(allreduce_hook)".
Example::
      >>> state = PowerSGDState(process_group=process_group, matrix_approximation_rank=1, start_powerSGD_iter=10)
      >>> ddp_model.register_comm_hook(state, fp16_compress_wrapper(powerSGD_hook))
Return type:
      Callable[[Any, GradBucket], Future[Tensor]]
torch.distributed.algorithms.ddp_comm_hooks.default_hooks.bf16_compress_wrapper(hook)
Warning: This API is experimental, and it requires NCCL version
   later than 2.9.6.
This wrapper casts the input gradient tensor of a given DDP
   communication hook to half-precision Brain floating point format
   https://en.wikipedia.org/wiki/Bfloat16_floating-point_format _
   (``torch.bfloat16), and casts the resulting tensor of the given
   hook back to the input data type, such as "float32".
Therefore, "bf16_compress_hook" is equivalent to
   "bf16_compress_wrapper(allreduce_hook)".
Example::
Example::
      >>> state = PowerSGDState(process_group=process_group, matrix_approximation_rank=1, start_powerSGD_iter=10)
      >>> ddp_model.register_comm_hook(state, bf16_compress_wrapper(powerSGD_hook))
Return type:
      Callable[[Any, GradBucket], Future[Tensor]]
PowerSGD Communication Hook
PowerSGD (Vogels et al., NeurIPS 2019) is a gradient compression
algorithm, which can provide very high compression rates and
accelerate bandwidth-bound distributed training. This algorithm needs
to maintain both some hyperparameters and the internal state.
Therefore, PowerSGD communication hook is a stateful hook, and the
user needs to provide a state object defined as below.
PowerSGD State
PowerSGD State
class torch.distributed.algorithms.ddp_comm_hooks.powerSGD_hook.PowerSGDState(process_group, matrix_approximation_rank=1, start_powerSGD_iter=1000, min_compression_rate=2, use_error_feedback=True, warm_start=True, orthogonalization_epsilon=0, random_seed=0, compression_stats_logging_frequency=10000, batch_tensors_with_same_shape=False)
Stores both the algorithm's hyperparameters and the internal state
   for all the gradients during the training. Particularly,
   "matrix_approximation_rank" and "start_powerSGD_iter" are the main
   hyperparameters that should be tuned by the user. For performance,
   we suggest to keep binary hyperparameters "use_error_feedback" and
   "warm_start" on.

"matrix_approximation_rank" controls the size of compressed low-
      rank tensors, which determines the compression rate. The lower
      the rank, the stronger the compression. 1.1. If "matrix_approximation_rank" is too low, the full
model quality will need more training steps to reach or will
         never reach and yield loss in accuracy.
     1.2. The increase of "matrix_approximation_rank" can
     substantially increase the computation costs of the
     compression, and the accuracy may not be further improved
     beyond a certain "matrix_approximation_rank" threshold.

To tune "matrix_approximation_rank", we suggest to start from 1 and
   increase by factors of 2 (like an exponential grid search, 1, 2, 4,
   ...), until a satisfactory accuracy is reached. Typically only a
   small value 1-4 is used. For some NLP tasks (as shown in Appendix D
   of the original paper), this value has been increased to 32.

"start_powerSGD_iter" defers PowerSGD compression until step
      "start_powerSGD_iter", and vanilla allreduce runs prior to step
      "start_powerSGD_iter". This hybrid scheme of **vanilla allreduce
PowerSGD** can effectively improve the accuracy, even a
relatively small "matrix_approximation_rank" is used. This is
      because that, the beginning of training phase is usually very
      sensitive to inaccurate gradients, and compressing gradients too
      early may make the training quickly take a suboptimal
      trajectory, which can result in an irrecoverable impact on the
      accuracy.
To tune "start_powerSGD_iter", we suggest to start with 10% of
   total training steps, and increase it until a satisfactory accuracy
   is reached. If there is a warm-up stage in the training,
   "start_powerSGD_iter" typically should be no less than the number
   of warm-up steps.

"min_compression_rate" is the minimum compression rate required
      when a layer is compressed. Due to the computation overheads
      incurred by the compression, a tensor is worth compressing only
      if there can be sufficient saving in bandwidth, where "(num_rows
num_cols) * matrix_approximation_rank * min_compression_rate <
num_rows * num_cols". If the specified compression rate
      threshold cannot be satisfied, the tensor will be directly
      allreduced without compression.
Compression statistics are logged every
   "compression_stats_logging_frequency" iterations once PowerSGD
   compression starts.


"orthogonalization_epsilon" can be a very small value (e.g.,
      1e-8) added to every normalized matrix column in
      orthogonalization step, to prevent div-by-zero error if any
      column has all 0s. If this can already be prevented (e.g., by
      batch normalization), an epsilon of 0 is recommended for
      accuracy.


"batch_tensors_with_same_shape" controls whether to compress and
      decompress tensors with same shape in a batched operation to
      achieve higher parallelism. Note that you should also increase
      the bucket size (i.e., "bucket_cap_mb" arg in DDP constructor)
      to make more same-shaped tensors appear in the same bucket,
however this may reduce the overlap between computation and
      communication, and increase the memory footprint due to stacking
      the tensors of the same shape. Set to "True" if the compression
      / decompression computation is a bottleneck.
Warning:
 If error feedback or warm-up is enabled, the minimum value of
 "start_powerSGD_iter" allowed in DDP is 2. This is because there
 is another internal optimization that rebuilds buckets at
 iteration 1 in DDP, and this can conflict with any tensor
 memorized before the rebuild process.

PowerSGD Hooks
Warning:
PowerSGD typically requires extra memory of the same size as the
  model's gradients to enable error feedback, which can compensate for
  biased compressed communication and improve accuracy.
Warning:
PowerSGD hooks may conflict with Apex automatic mixed precision
  package. Please use PyTorch native automatic mixed precision package
  instead.
instead.
torch.distributed.algorithms.ddp_comm_hooks.powerSGD_hook.powerSGD_hook(state, bucket)
This DDP communication hook implements PowerSGD gradient
   compression algorithm described in the paper. Once gradient tensors
   are aggregated across all workers, this hook applies compression as
   follows:


Views the input flattened 1D gradient tensor as a list of per-
      parameter tensors, and divides all the tensors into two groups:
 1.1 The tensors that should be compressed before allreduce,
 because the compression can give enough saving in bandwidth.

 1.2 Rest of the tensors will be directly allreduced without
 compression, including all the vector tensors (for biases).



Handles uncompressed tensors:
 2.1. Allocate contiguous memory for those uncompressed
 tensors, and allreduces all the uncompressed tensors as a
 batch, without compression;

 2.2. Copies the individual uncompressed tensors from the
contiguous memory back to the input tensor.

Handles the tensors that should be compressed by PowerSGD
      compression: 3.1. For each tensor M, creates two low-rank tensors P and Q
 for decomposing M, such that M = PQ^T, where Q is initialized
 from a standard normal distribution and orthogonalized;

 3.2. Computes each P in Ps, which is equal to MQ;

 3.3. Allreduces Ps as a batch;

 3.4. Orthogonalizes each P in Ps;

 3.5. Computes each Q in Qs, which is approximately equal to
 M^TP;

 3.6. Allreduces Qs as a batch;

 3.7. Computes each M among all the compressed tensors, which
 is approximately equal to PQ^T.



Note that this communication hook enforces vanilla allreduce for
   the first "state.start_powerSGD_iter" iterations. This not only
   gives the user more control over the tradeoff between speedup and
   accuracy, but also helps abstract away some complexity of the