Skip to content

Commit

Permalink
Update doc from commit e8720ba
Browse files Browse the repository at this point in the history
  • Loading branch information
@torchxlabot2
torchxlabot2 committed Dec 10, 2024
1 parent 40b58c6 commit aa78c1a
Show file tree
Hide file tree
Showing 45 changed files with 56 additions and 44 deletions.
2 changes: 1 addition & 1 deletion master/_modules/index.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -265,7 +265,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/_modules/torch_xla/core/xla_model.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -265,7 +265,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/_modules/torch_xla/debug/metrics.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -265,7 +265,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/_modules/torch_xla/distributed/parallel_loader.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -265,7 +265,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
14 changes: 13 additions & 1 deletion master/_modules/torch_xla/distributed/spmd/xla_sharding.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -265,7 +265,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down Expand Up @@ -427,6 +427,8 @@ <h1>Source code for torch_xla.distributed.spmd.xla_sharding</h1><div class="high
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Tuple</span><span class="p">,</span> <span class="n">Union</span><span class="p">,</span> <span class="n">List</span><span class="p">,</span> <span class="n">Sequence</span><span class="p">,</span> <span class="n">Any</span><span class="p">,</span> <span class="n">Optional</span><span class="p">,</span> <span class="n">Set</span>
<span class="kn">from</span> <span class="nn">enum</span> <span class="kn">import</span> <span class="n">IntEnum</span>

<span class="kn">from</span> <span class="nn">torch.amp</span> <span class="kn">import</span> <span class="n">custom_fwd</span><span class="p">,</span> <span class="n">custom_bwd</span>


<div class="viewcode-block" id="Mesh"><a class="viewcode-back" href="../../../../learn/api-guide.html#torch_xla.distributed.spmd.Mesh">[docs]</a><span class="k">class</span> <span class="nc">Mesh</span><span class="p">:</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Describe the logical XLA device topology mesh and the underlying resources.</span>
Expand Down Expand Up @@ -1147,10 +1149,19 @@ <h1>Source code for torch_xla.distributed.spmd.xla_sharding</h1><div class="high
<span class="sd"> dimensions. The torch.matmul default behavior makes it very hard for XLA compiler</span>
<span class="sd"> to propagate the sharding annotation.</span>

<span class="sd"> Autocast decorators @custom_fwd and @custom_bwd used as per autocast docs [1] to bring this class/layer within </span>
<span class="sd"> autocast context, when autocast is enabled.</span>
<span class="sd"> torch.get_autocast_dtype() fetches datatype for ops run in autocast [2], with the specified device (here, &#39;xla&#39;).</span>
<span class="sd"> </span>
<span class="sd"> References: </span>
<span class="sd"> [1] https://pytorch.org/docs/stable/notes/amp_examples.html#functions-with-multiple-inputs-or-autocastable-ops </span>
<span class="sd"> [2] https://github.com/pytorch/pytorch/blob/2cc01cc6d3ad2aff47e8460667ba654b2e4c9f21/torch/amp/autocast_mode.py#L500</span>

<span class="sd"> TODO (alanwaketan): Let&#39;s patch it on the dispatcher level.</span>
<span class="sd"> &quot;&quot;&quot;</span>

<span class="nd">@staticmethod</span>
<span class="nd">@custom_fwd</span><span class="p">(</span><span class="n">device_type</span><span class="o">=</span><span class="s1">&#39;xla&#39;</span><span class="p">,</span> <span class="n">cast_inputs</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">get_autocast_dtype</span><span class="p">(</span><span class="s1">&#39;xla&#39;</span><span class="p">))</span>
<span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="nb">input</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">bias</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
<span class="c1"># bias is an optional argument</span>
<span class="n">ctx</span><span class="o">.</span><span class="n">save_for_backward</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">bias</span><span class="p">)</span>
Expand All @@ -1161,6 +1172,7 @@ <h1>Source code for torch_xla.distributed.spmd.xla_sharding</h1><div class="high
<span class="k">return</span> <span class="n">product</span> <span class="o">+</span> <span class="n">bias</span>

<span class="nd">@staticmethod</span>
<span class="nd">@custom_bwd</span><span class="p">(</span><span class="n">device_type</span><span class="o">=</span><span class="s1">&#39;xla&#39;</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">backward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">grad_output</span><span class="p">):</span>
<span class="nb">input</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">bias</span> <span class="o">=</span> <span class="n">ctx</span><span class="o">.</span><span class="n">saved_tensors</span>
<span class="n">grad_input</span> <span class="o">=</span> <span class="n">grad_weight</span> <span class="o">=</span> <span class="n">grad_bias</span> <span class="o">=</span> <span class="kc">None</span>
Expand Down
2 changes: 1 addition & 1 deletion master/_modules/torch_xla/distributed/xla_multiprocessing.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -265,7 +265,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/_modules/torch_xla/experimental/eager.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -265,7 +265,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/_modules/torch_xla/runtime.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -265,7 +265,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/_modules/torch_xla/torch_xla.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -265,7 +265,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/accelerators/gpu.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/accelerators/tpu.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/contribute/bazel.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -266,7 +266,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/contribute/codegen_migration.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/contribute/configure-environment.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/contribute/op_lowering.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/contribute/plugins.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/features/distop.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -265,7 +265,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/features/pallas.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/features/stablehlo.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/features/triton.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/genindex.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -265,7 +265,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/index.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -266,7 +266,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/learn/api-guide.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/learn/dynamic_shape.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/learn/eager.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/learn/pjrt.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/learn/pytorch-on-xla-devices.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/learn/troubleshoot.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/learn/xla-overview.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/notes/source_of_recompilation.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -265,7 +265,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
Binary file modified master/objects.inv
View file
Binary file not shown.
2 changes: 1 addition & 1 deletion master/perf/amp.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/perf/ddp.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/perf/dynamo.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/perf/fori_loop.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/perf/fsdp.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/perf/fsdpv2.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/perf/quantized_ops.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/perf/recompilation.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/perf/spmd_advanced.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
2 changes: 1 addition & 1 deletion master/perf/spmd_basic.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@


<div class="version">
master (2.6.0+git1812817 )
master (2.6.0+gite8720ba )
</div>


Expand Down
Loading

0 comments on commit aa78c1a

Please sign in to comment.