haliax 1.4.dev332__tar.gz → 1.4.dev333__tar.gz

{haliax-1.4.dev332 → haliax-1.4.dev333}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev332
+Version: 1.4.dev333
 Summary: Named Tensors for Legible Deep Learning in JAX
 Project-URL: Homepage, https://github.com/stanford-crfm/haliax
 Project-URL: Bug Tracker, https://github.com/stanford-crfm/haliax/issues/

{haliax-1.4.dev332 → haliax-1.4.dev333}/docs/api.md

@@ -55,11 +55,11 @@ Occasionally, an axis size can be inferred in some circumstances but not others.
 
 ### Combining Arrays
 
-We don't include `hstack` or `vstack` because we prefer semantic axes.
 
 ::: haliax.concatenate
 ::: haliax.stack
 
+(We don't include `hstack` or `vstack` because they are subsumed by `stack`.)
 
 ## Array Manipulation
 
@@ -289,6 +289,16 @@ See also the section on [Partitioning](partitioning.md).
 ::: haliax.partitioning.physical_axis_size
 ::: haliax.partitioning.sharding_for_axis
 
+
+## Gradient Checkpointing
+
+Haliax mainly just defers to JAX and [equinox.filter_checkpoint][] for gradient checkpointing. However,
+we provide a few utilities to make it easier to use.
+
+See also [haliax.nn.StackedCheckpointPolicy][].
+
+::: haliax.tree_checkpoint_name
+
 ### Old API
 
 These functions are being deprecated and will be removed in a future release.

{haliax-1.4.dev332 → haliax-1.4.dev333}/docs/broadcasting.md

@@ -62,8 +62,8 @@ This method takes one or more axes and adds them to the array.
 
 ## Explicit Broadcasting Functions
 
-::: haliax.broadcast_axis
-::: haliax.broadcast_to
+* [haliax.broadcast_axis][]
+* [haliax.broadcast_to][]
 
 
 <!--

{haliax-1.4.dev332 → haliax-1.4.dev333}/docs/css/mkdocstrings.css

@@ -1,9 +1,19 @@
 /* Indentation. */
-div.doc-contents:not(.first) {
+div.doc-contents {
   padding-left: 25px;
   border-left: .05rem solid var(--md-typeset-table-color);
 }
 
+
+div.doc-class:not(.doc-contents .doc-contents)::after {
+  content: "";
+  display: block;
+  width: 100%;
+  height: 1px; /* Adjust thickness */
+  background-color: black; /* Adjust color */
+  margin: 10px 0; /* Adjust spacing */
+}
+
 /* Mark external links as such. */
 a.external::after,
 a.autorefs-external::after {

{haliax-1.4.dev332 → haliax-1.4.dev333}/docs/fp8.md

@@ -1,20 +1,11 @@
-# FP8 Training
+# Quantized Training
 
 !!! warning
 
-    FP8 training in Haliax is currently experimental and may change in the future.
+    FP8 and Int8 training in Haliax is currently experimental and may change in the future.
 
-FP8 refers to 8-bit floating point numbers. FP8 is a massively reduced precision compared to the 32-bit floating point numbers
-or 16-bit floating point numbers that are typically used in deep learning: there are only 256 possible values in FP8, compared to
-the (almost) 2^32 in 32-bit and 2^16 in 16-bit. However, FP8 is still useful for training deep learning models, especially on
-hardware that is optimized for FP8. In particular, it can massively accelerate training on hardware that is optimized for FP8:
-H100 has 2x FP8 FLOPS compared to FP16 FLOPS and almost 60x(!) compared to F32 FLOPS.
-
-The FP8 in Haliax is currently designed to optimize throughput on FP8-enabled devices (currently H100) rather
-than to save memory. In particular, Haliax's FP8 support is not designed to quantize a model to FP8 for deployment,
-though this shouldn't be that hard to add for models that were trained using this functionality.
-We would be happy to accept contributions to add this functionality,
-and are happy to work with you to do so. In particular, adding this for models trained using Haliax's FP8 should be easy.
+Haliax supports training with FP8 and int8. This is useful for training on hardware that is optimized for FP8 or Int8,
+such as the H100 (fp8) or A100s (int8) and TPU v5 and newer (int8).
 
 ## TL;DR
 
@@ -38,15 +29,42 @@ module = haxq.apply_updates(module, updates, overwrite)
 
 And train your model like normal.
 
-## How to use FP8
+Similarly, you can use `Int8` by setting `Int8=True` in the `QuantizationConfig` object.
 
-To use FP8, you need to do two things:
 
-* Enable FP8 for the layers you want to use FP8
-* Modify your training step to be compatible with FP8
 
-Each of these is just a couple of lines of code.
+## What is FP8?
+
+FP8 refers to 8-bit floating point numbers. FP8 is a massively reduced precision compared to the 32-bit floating point numbers
+or 16-bit floating point numbers that are typically used in deep learning: there are only 256 possible values in FP8, compared to
+the (almost) 2^32 in 32-bit and 2^16 in 16-bit. However, FP8 is still useful for training deep learning models, especially on
+hardware that is optimized for FP8. In particular, it can massively accelerate training on hardware that is optimized for FP8:
+H100 has 2x FP8 FLOPS compared to FP16 FLOPS and almost 60x(!) compared to F32 FLOPS.
+
+The FP8 in Haliax is currently designed to optimize throughput on FP8-enabled devices (currently H100) rather
+than to save memory. In particular, Haliax's FP8 support is not designed to quantize a model to FP8 for deployment,
+though this shouldn't be that hard to add for models that were trained using this functionality.
+We would be happy to accept contributions to add this functionality,
+and are happy to work with you to do so. In particular, adding this for models trained using Haliax's FP8 should be easy.
+
+See this [FP8 Primer](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/fp8_primer.html) for more information on FP8.
+
+## What is Int8?
 
+Int8 refers to 8-bit integers. Int8 has the same number of bits as FP8, but the interpretation is different: instead of
+exponentially spaced numbers, Int8 has linearly spaced numbers.
+
+In Haliax, we support Int8 training through Google's [AQT](https://github.com/google/aqt) library. AQT (for
+"Accurate Quantization Training") is a library that allows you to train models with quantization-aware training (QAT).
+
+## How to use FP8 or Int8 in Haliax
+
+To use quantized training in Haliax, you need to do three things:
+
+* Enable FP8 (or int8) for the layers you want
+* Modify your training step to be compatible
+
+Each of these is just a couple of lines of code.
 
 ```python
 import haliax as hax
@@ -96,8 +114,8 @@ updates, opt_state = opt.update(grads, opt_state, params=module)  # or however y
 module = hax.quantization.apply_updates(module, updates, grads)
 ```
 
-That's it! Just a few lines of code to enable FP8. The `quantize_linear_layers` function will transform your module to use FP8
-for linear layers (or a subset if you want), and the combo of `partition_for_grad_overwrite` and `apply_updates` function will apply the updates to the module
+That's it! Just a few lines of code to enable FP8. The `quantize_linear_layers` function will transform your module to use
+quantization-aware training for linear layers (or a subset if you want), and the combo of `partition_for_grad_overwrite` and `apply_updates` function will apply the updates to the module
 in a way that is compatible with FP8.
 
 ## How FP8 works
@@ -134,6 +152,13 @@ original precision.  It remembers the maximum absolute value for each of the inp
 and scales the gradients back to the original precision. It remembers the maximum absolute value for the incoming
 gradient and stores it in the gradient.
 
+## How Int8 works
+
+Int8 is in principle the same, though the details differ. AQT is a much more flexible library than the FP8 implementation,
+because it can be a bit more finicky. We use AQT directly, and we recommend you look at the
+[AQT documentation](https://github.com/google/aqt?tab=readme-ov-file#how-aqt-works-internally) for more
+information on how it works.
+
 # API Reference
 
 ## Functions
@@ -152,7 +177,8 @@ gradient and stores it in the gradient.
 
 ::: haliax.quantization.DefaultDotGeneralOp
 ::: haliax.quantization.Fp8DotGeneralOp
+::: haliax.quantization.Int8DotGeneralOp
 
 ## Configuration
 
-::: haliax.quantization.Fp8Config
+::: haliax.quantization.QuantizationConfig

{haliax-1.4.dev332 → haliax-1.4.dev333}/docs/nn.md

@@ -5,6 +5,7 @@
 
 Haliax provides a small number of neural network modules that are compatible with Equinox, though
 they naturally all use [haliax.NamedArray][]. (We welcome PRs for more modules! Nothing too exotic though.)
+
 The most interesting of these modules is [haliax.nn.Stacked][], which allows you to create homogenous "stacks"
 of the same module (e.g. transformer blocks), which is a common pattern in deep learning.
 
@@ -21,11 +22,12 @@ of the same module (e.g. transformer blocks), which is a common pattern in deep
 
 ### Meta
 
-::: haliax.nn.scan.BlockFoldable
-::: haliax.nn.BlockSeq
-::: haliax.nn.Stacked
 ::: haliax.nn.MLP
 
+### Stacked
+
+See the full documentation of [Stacked](stacked.md).
+
 ### Convolution
 
 Unlike other frameworks, Haliax doesn't distinguish between 1D, 2D, and 3D, and general convolutions. Instead, we have

haliax-1.4.dev333/docs/stacked.md

@@ -0,0 +1,203 @@
+# Module Stacks
+
+A core pattern for larger models in JAX is the "scan-over-layers" pattern, where you have a sequence of layers
+that get stacked together, and you use [jax.lax.scan][] or [haliax.fold][] or [haliax.scan][] to apply them to a
+sequence of inputs. In Haliax, layers are represented as [equinox.nn.Module][]s, and the [haliax.nn.Stacked][] module
+provides a way to create a sequence of layers that can be applied to a sequence of inputs that implements the
+scan-over-layers pattern.
+
+## Stacked
+
+[haliax.nn.Stacked][] lets you apply a layer sequentially to an input, scanning over a "Layers" axis. For instance,
+a Transformer might use a Stacked for its Transformer blocks:
+
+
+```python
+class TransformerBlock(eqx.Module):
+
+    def __init__(self, config: TransformerConfig, layer_index, *, key):
+        attn_key, mlp_key = jax.random.split(key)
+        self.attention = Attention.init(config, key=attn_key)
+        self.mlp = MLP.init(config, key=mlp_key)
+        self.ln1 = LayerNorm.init(config.Hidden)
+        self.ln2 = LayerNorm.init(config.Hidden)
+        self.layer_index = layer_index
+
+    def __call__(self, x):
+        y = self.attention(self.ln1(x))
+        x = x + y
+        y = self.mlp(self.ln2(x))
+        return x + y
+
+class Transformer(eqx.Module):
+    def __init__(self, config: TransformerConfig):
+        self.blocks = Stacked.init(Layers, TransformerBlock)(
+            config,  # static configuration
+            scale=hax.arange(Layers),  # dynamic configuration. Each layer gets a scalar scale value [0, 1, 2, ...]
+            key=jax.random.split(key, Layers),  # dynamic configuration. Each layer gets a different key
+        )
+
+    def __call__(self, x: NamedArray) -> NamedArray:
+        # morally the equivalent of:
+        # for block in self.blocks:
+        #     x = block(x)
+        # Except that it works better with JAX compile times.
+
+        return self.blocks.fold(x)
+```
+
+You can think of [haliax.nn.Stacked][] as an analog to PyTorch's
+[torch.nn.Sequential](https://pytorch.org/docs/stable/generated/torch.nn.Sequential.html), except that
+every layer in the sequence must have exactly the same shape and configuration.
+
+Internally, a Stacked is a single copy of the module, except that every NamedArray inside that module
+has a Block axis prepended (as though they were stacked with [haliax.stack][]). Similarly, every JAX array
+inside the module has its first axis prepended with an axis of the same size as the Block axis, as though
+they were stacked with [jax.numpy.stack][].
+
+When you call the Stacked, it scans over the Block axis, applying the module to each element of the Block.
+
+### Creating a Stacked
+
+To create a Stacked, we provide `Stacked.init`, which takes a "Layers" [haliax.Axis][] and another Module as
+well as args and kwargs for that module. The Layer is the axis that the Stacked will scan over, and the `args`
+and `kwargs` are implicitly vmapped over the Layers.
+
+For instance, to create a stack of GPT2 blocks, you might do:
+
+```python
+import jax.random
+
+blocks = Stacked.init(Layers, Gpt2Block)(
+    config,  # static configuration
+    scale=hax.arange(Layers),  # dynamic configuration. Each layer gets a scalar scale value [0, 1, 2, ...]
+    key=jax.random.split(key, Layers.size),  # dynamic configuration. Each layer gets a different key
+)
+```
+
+Any NamedArray passed to the Stacked init will have its Layers axis (if present) vmapped over. Any
+JAX array will have its first axis vmapped over.
+
+
+### Fold Blocks vs Scan Blocks
+
+The Stacked module provides two ways to apply the layers: `fold` and `scan`.  A fold is the moral equivalent of this for loop:
+
+```python
+for block in self.blocks:
+    x = block(x)
+```
+
+while a scan is the moral equivalent of this for loop:
+
+```python
+out = []
+for block in self.blocks:
+    x, y = block(x)
+    out.append(y)
+
+return x, stack(out)
+```
+
+Blocks can be coded to either support fold or scan, but not both.
+A "fold Block" should have the signature `def __call__(self, x: Carry) -> Carry`,
+while a "scan Block" should have the signature `def __call__(self, x: Carry) -> Tuple[Carry, Output]`.
+
+(See also [jax.lax.scan][], [haliax.fold][], and [haliax.scan][].)
+
+### Requirements for Stacked Blocks
+
+As we said above, the Stacked module requires that all the layers have the same shape and configuration.
+
+A further constraint is that the elements of the stack must have the same Python control flow. This is the usual
+constraint imposed on jit-compiled functions in JAX. All control flow must use `jax.lax` primitives like
+[jax.lax.cond][], [jax.lax.while_loop][], and [jax.lax.scan][]. You can't use Python control flow like `if` or `for`
+except for static control flow that is the same for all elements of the stack.
+
+## BlockSeq and BlockFoldable
+
+We also provide a way to create a sequence of layers that can be applied to a sequence of inputs that implements the
+same interface as [haliax.nn.Stacked][], but with a different implementation. This is the [haliax.nn.BlockSeq][] module.
+BlockSeq implements those for loops directly, rather than using [haliax.fold][] or [haliax.scan][].
+
+[haliax.nn.scan.BlockFoldable][] is an interface that both [haliax.nn.Stacked][] and [haliax.nn.BlockSeq][] implement.
+
+## Gradient Checkpointing
+
+The [haliax.nn.Stacked][] module also provides a way to do gradient checkpointing, which can be useful for deep models.
+
+Gradient checkpointing, aka rematerialization, is a technique for trading off memory usage for compute time.
+Instead of storing all the intermediate activations of a model, you store only a subset and recompute the rest
+as needed. (XLA automatically recomputes the rest for you as needed.)
+
+[JAX's checkpointing mechanism]((https://docs.jax.dev/en/latest/gradient-checkpointing.html) is highly flexible,
+and we provide a relatively simple interface to it for use with `Stacked`.
+
+### Simple Checkpointing
+In the simplest case, you can enable a usually-good-enough checkpointing policy by passing `gradient_checkpointing=True`
+to the `Stacked.init` call:
+
+```python
+blocks = Stacked.init(Layers, TransformerBlock, gradient_checkpointing=True)(
+    config,
+    scale=hax.arange(Layers),
+    key=jax.random.split(key, Layers.size),
+)
+```
+
+This will preserve the intermediate "carries" and the "outputs" of the scans, while rematerializing (i.e. recomputing)
+the rest of the computation as needed during backpropagation.
+
+### Custom Checkpointing Policies
+
+If you need more control over the checkpointing policy, you can pass a [haliax.nn.StackedCheckpointPolicy][] object to
+the Stacked init:
+
+```python
+policy = StackedCheckpointPolicy(
+   save_carries=True,  # default
+   save_outputs=True,  # default
+   save_intermediates=False,  # default
+)
+```
+
+### Saving Block-Internal Values
+
+"`intermediates`" refers to the internal computation of the block. If you set `save_intermediates=True`, then
+all internals of every block will be saved. This can be expensive.
+
+You can also pass a list of strings to `save_intermediates` to specify which intermediates to save.
+
+You could, for instance, save the output of the attention layer using [haliax.tree_checkpoint_name][]:
+
+```python
+class TransformerBlock(eqx.Module):
+    def __call__(self, x):
+        y = self.attention(self.ln1(x))
+        y = haliax.tree_checkpoint_name(y, "attn_out")
+        x = x + y
+        y = self.mlp(self.ln2(x))
+        return x + y
+
+policy = StackedCheckpointPolicy(save_carries=True, save_block_internals=["attn_out"])
+```
+
+With this policy, the output of the attention layer will be saved during the forward pass.
+
+This will save an extra attention computation in the backward pass, adding $`O(N * Pos * Hidden)`$ memory usage,
+which is double that required by the default policy.
+
+
+### Offloading Checkpointed Values
+
+Both `save_carries` and `save_outputs` can either be a boolean or the string "offload". If "offload", then the
+checkpointed values will be offloaded to the host during the forward pass, and reloaded during the backward pass.
+
+
+## API
+
+::: haliax.nn.Stacked
+::: haliax.nn.BlockSeq
+::: haliax.nn.scan.BlockFoldable
+
+::: haliax.nn.StackedCheckpointPolicy

{haliax-1.4.dev332 → haliax-1.4.dev333}/docs/state-dict.md

@@ -157,4 +157,21 @@ to join the prefix to the keys of the state dict.
 
 ## API Reference
 
-::: haliax.state_dict
+::: haliax.state_dict.StateDict
+::: haliax.state_dict.ModuleWithStateDictSerialization
+
+### Saving and Loading State Dicts
+::: haliax.state_dict.save_state_dict
+::: haliax.state_dict.load_state_dict
+
+### Converting betweewn State Dicts and Modules
+
+::: haliax.state_dict.from_state_dict
+::: haliax.state_dict.to_state_dict
+
+### Torch Compatibility
+
+::: haliax.state_dict.from_torch_compatible_state_dict
+::: haliax.state_dict.to_torch_compatible_state_dict
+::: haliax.state_dict.flatten_linear_layers
+::: haliax.state_dict.unflatten_linear_layers

{haliax-1.4.dev332 → haliax-1.4.dev333}/mkdocs.yml

@@ -43,25 +43,34 @@ plugins:
             - https://jax.readthedocs.io/en/latest/objects.inv
             - https://docs.kidger.site/equinox/objects.inv
             - https://einops.rocks/objects.inv
+          selection:
+            inherited_members: true
           options:
             docstring_options:
               ignore_init_summary: true
               docstring_style: sphinx
             show_source: false
-            docstring_section_style: list
-            heading_level: 5
+            filters:
+              - "!^_"
+            heading_level: 4
             inherited_members: true
+            members_order: source
             merge_init_into_class: true
+            parameter_headings: true
+            separate_signature: false
             load_external_modules: true
             preload_modules: [haliax, haliax.core]
-#            separate_signature: true
+            show_if_no_docstring: true
             show_root_heading: true
             show_root_full_path: false
-#            show_signature_annotations: true
-            show_symbol_type_heading: false
+            show_signature_annotations: true
+            docstring_section_style: list
+            show_symbol_type_heading: true
             show_symbol_type_toc: false
             signature_crossrefs: true
             line_length: 100
+            summary: true
+
   - include-markdown
 extra_css:
   - css/material.css

haliax-1.4.dev333/src/haliax/__about__.py

@@ -0,0 +1 @@
+__version__ = "1.4.dev333"

{haliax-1.4.dev332 → haliax-1.4.dev333}/src/haliax/_src/state_dict.py

@@ -395,7 +395,7 @@ def flatten_linear_layers(tree: T) -> T:
         new_In: Axis = flatten_axes(layer.In, "__IN__")
 
         if weight.array is not None:
-            out_first = layer.out_first
+            out_first = layer._out_first
             weight = weight.flatten_axes(layer.Out, new_Out).flatten_axes(layer.In, new_In)
 
             if out_first:

{haliax-1.4.dev332 → haliax-1.4.dev333}/src/haliax/axis.py

@@ -35,6 +35,8 @@ def make_axes(**kwargs: int) -> Tuple[Axis, ...]:
     Example:
     ```
     X, Y = axes(X=10, Y=20)
+    ```
+
     """
     return tuple(Axis(name, size) for name, size in kwargs.items())
 

{haliax-1.4.dev332 → haliax-1.4.dev333}/src/haliax/nn/embedding.py

@@ -24,26 +24,60 @@ class Embedding(eqx.Module):
 
     @staticmethod
     def init(Vocab: Axis, Embed: AxisSpec, *, init_scale: float = 1, key, initializer_range: Optional[float] = None):
+        """
+        Initialize an Embedding module.
+
+        An embedding module is a simple lookup table that maps integer indices to vectors or tensors.
+        Weights are initialized with a truncated normal distribution with a standard deviation of
+          `init_scale / output_size`.
+
+        Args:
+            Vocab: Size of the vocabulary
+            Embed: Shape of the embedding vectors. May be a single axis or a full AxisSpec
+            init_scale: Scale of the initialization
+            key: PRNG key
+            initializer_range: Deprecated. Use init_scale instead.
+        """
         if initializer_range is not None:
             warnings.warn("initializer_range is deprecated. Use init_std instead.", DeprecationWarning)
             init_scale = initializer_range
 
         all_axes = (Vocab,) + ensure_tuple(Embed)
         output_size = hax.axis_size(Embed)
-        weight = hax.random.truncated_normal(key, all_axes, -3, 3) * (init_scale / math.sqrt(output_size))
+        weight = hax.random.truncated_normal(key, all_axes, -3, 3) * (init_scale / output_size)
         return Embedding(weight=weight, Vocab=Vocab, Embed=Embed)
 
-    def __call__(self, input_ids, *, key: Optional[PRNGKeyArray] = None):
+    def __call__(self, input_ids: NamedArray, *, key: Optional[PRNGKeyArray] = None):
+        """Alias for `embed`. key is ignored."""
         return self.embed(input_ids)
 
     @named_call
-    def embed(self, input_ids):
+    def embed(self, input_ids: NamedArray):
+        """
+        Args:
+            input_ids: token IDs with shape > {Vocab}
+        """
         input_embeds = self.weight.take(self.Vocab, input_ids)
         return input_embeds
 
-    def unembed(self, input_embeds):
+    def unembed(self, input_embeds: NamedArray):
+        """
+        Unembed the input embeddings back to the vocabulary space.
+
+        Equivalent to `input_embeds.dot(self.weight, axis=self.Embed)`.
+        """
         return input_embeds.dot(self.weight, axis=self.Embed)
 
     def resize_embeddings(self, new_size: int, key: Optional[PRNGKeyArray] = None):
+        """
+        Resize the embedding layer to a new size.
+        Args:
+            new_size: New size of the vocabulary
+            key: PRNG key for initialization of any new weights
+
+        Returns:
+            Embedding: Resized embedding layer
+
+        """
         new_weights = resize_axis(self.weight, self.Vocab, new_size, key=key)
         return dataclasses.replace(self, Vocab=self.Vocab.resize(new_size), weight=new_weights)  # type: ignore

{haliax-1.4.dev332 → haliax-1.4.dev333}/src/haliax/nn/linear.py

@@ -72,7 +72,10 @@ class Linear(eqx.Module):
         return q
 
     @property
-    def out_first(self):
+    def _out_first(self):
+        """
+        Returns: bool: Whether the output axes are first in the weight matrix
+        """
         # We do it this way because of scan layers
         if isinstance(self.Out, hax.Axis):
             return self.weight.axes[-1] != self.Out

{haliax-1.4.dev332 → haliax-1.4.dev333}/src/haliax/nn/normalization.py

@@ -58,6 +58,9 @@ def logsumexp(a: A, axis: Optional[AxisSelection] = None) -> A:
     return wrap_reduction_call(jnn.logsumexp, a, axis=axis, single_axis_only=False, supports_where=False)
 
 
+# TODO: support where in softmax, etc
+
+
 def softmax(a: A, axis: Optional[AxisSelection] = None) -> A:
     return wrap_axiswise_call(jnn.softmax, a, axis=axis, single_axis_only=False)
 

haliax-1.4.dev332/src/haliax/__about__.py

@@ -1 +0,0 @@
-__version__ = "1.4.dev332"