Introduce float8 training

[james77777778]

This PR introduces float8 training, especially for TF and JAX.

Notes

A New Attribute overwrite_with_gradient Has Been Added to Variable

This attribute will be utilized in Optimizer if set to True. The optimizer will directly overwrite the variable by the computed gradient in BaseOptimizer._overwrite_variables_directly_with_gradients of BaseOptimizer.apply.

With this new behavior, we can correctly update float8 parameters such as scale and amax_history.

A New Quantization Mode float8 Has Been Added

To DTypePolicy:

# This should be only invoked by quantizable layers.
original_policy = dtype_policy().name
policy = QuantizedDTypePolicy(f"float8_from_{original_policy}")

To Dense and EinsumDense:

# Dense
model = models.Sequential([layers.Dense(units=16)])
model.quantize("float8")
model.compile(optimizer="sgd", loss="mse")
model.fit(x, y, epochs=2)

# EinsumDense
model = models.Sequential([layers.EinsumDense("ab,bcd->acd", (8, 32))])
model.quantize("float8")
model.compile(optimizer="sgd", loss="mse")
model.fit(x, y, epochs=2)

The compiled graph in TF should contains the expected op __cublas$lt$matmul$f8:

generate_hlo_dot.py

import tensorflow as tf
from keras import backend
from keras import config
from keras import layers
from keras import ops

keras.config.set_dtype_policy("mixed_bfloat16")

layer = layers.EinsumDense("abcde,aebf->adbcf", output_shape=[16, 16, 16, 16])
layer.build([32, 16, 16, 16, 16])
x = ops.ones([32, 16, 16, 16, 16])
layer.quantize("float8")

def fn(x):
    return layer(x)

compiled_fn = tf.function(fn, jit_compile=True)
y = compiled_fn(x)
ir = compiled_fn.experimental_get_compiler_ir(x)(stage="optimized_hlo_dot")
with open("ir_tf.dot", "w") as f:
    f.write(ir)

Results

• RTX 4070 12GB, CUDA12.3, CUDNN8.9
• GPT2 Base: batch_size=64, frozen embeddings
• GPT2 Medium: batch_size=24, frozen embeddings

Backend	Model	Recipe	Training Speed	Peak Mem. Usage	Val. Acc.
TF	GPT2 Base	mixed_bfloat16	274~275 ms/step	6.234GB	0.4997
TF	GPT2 Base	float8	269~272 ms/step	5.909GB	0.4971
TF	GPT2 Medium	mixed_bfloat16	252~254 ms/step	7.338GB	0.5535
TF	GPT2 Medium	float8	268~271 ms/step	7.002GB	0.5513
JAX	GPT2 Base	mixed_bfloat16	207~208 ms/step	X	0.4991
JAX	GPT2 Base	float8	200~202 ms/step	X	0.4977
JAX	GPT2 Medium	mixed_bfloat16	236~238 ms/step	X	0.5534
JAX	GPT2 Medium	float8	232~234 ms/step	X	0.5521

Standalone script:

fp8.py

"""
Dataset:
https://huggingface.co/datasets/databricks/databricks-dolly-15k/blob/main/databricks-dolly-15k.jsonl
"""

import argparse
import json

import kagglehub
import keras_nlp

import keras


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model",
        default="gpt2_base_en",
        choices=["gpt2_base_en", "gpt2_medium_en", "gemma_2b_en"],
        help="Which model to demonstrate",
    )
    parser.add_argument(
        "--fp8",
        action="store_true",
        help="Whether to use float8 technique",
    )
    parser.add_argument("-b", "--batch-size", default=32, type=int)
    args = parser.parse_args()
    return args


def get_optimizer_and_loss():
    optimizer = keras.optimizers.AdamW(
        learning_rate=5e-5,
        weight_decay=0.01,
        epsilon=1e-6,
        global_clipnorm=1.0,  # Gradient clipping.
    )
    # Exclude layernorm and bias terms from weight decay.
    optimizer.exclude_from_weight_decay(var_names=["bias"])
    optimizer.exclude_from_weight_decay(var_names=["gamma"])
    optimizer.exclude_from_weight_decay(var_names=["beta"])

    loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
    return optimizer, loss


class GPUMemoryCallback(keras.callbacks.Callback):
    def __init__(self, target_batches, **kwargs):
        super().__init__(**kwargs)
        self.target_batches = target_batches
        self.memory_usage = []

    def _compute_memory_usage(self):
        if keras.backend.backend() == "tensorflow":
            import tensorflow as tf

            memory_stats = tf.config.experimental.get_memory_info("GPU:0")
        else:
            memory_stats = {"peak": 0.0}
        # Convert bytes to GB and store in list.
        peak_usage = round(memory_stats["peak"] / (2**30), 3)
        self.memory_usage.append(peak_usage)

    def on_epoch_begin(self, epoch, logs=None):
        self._compute_memory_usage()

    def on_train_batch_begin(self, batch, logs=None):
        if batch in self.target_batches:
            self._compute_memory_usage()

    def on_epoch_end(self, epoch, logs=None):
        self._compute_memory_usage()


if __name__ == "__main__":
    EPOCHS = 2
    keras.config.disable_traceback_filtering()
    keras.mixed_precision.set_global_policy("mixed_bfloat16")

    args = get_args()
    if args.model == "gemma_2b_en":
        kagglehub.login()

    # Setup dataset
    data = []
    with open("databricks-dolly-15k.jsonl") as file:
        for line in file:
            features = json.loads(line)
            # Filter out examples with context, to keep it simple.
            if features["context"]:
                continue
            # Format the entire example as a single string.
            template = "Instruction:\n{instruction}\n\nResponse:\n{response}"
            data.append(template.format(**features))

    # Only use 16000 training examples, to keep it fast.
    data = data[:16000]

    if args.model == "gemma_2b_en":
        preprocessor = keras_nlp.models.GemmaCausalLMPreprocessor.from_preset(
            args.model, sequence_length=128
        )
        model = keras_nlp.models.GemmaCausalLM.from_preset(
            args.model, preprocessor=preprocessor
        )
        model.backbone.token_embedding.trainable = False
    elif "gpt2" in args.model:
        preprocessor = keras_nlp.models.GPT2CausalLMPreprocessor.from_preset(
            args.model, sequence_length=128
        )
        model = keras_nlp.models.GPT2CausalLM.from_preset(
            args.model, preprocessor=preprocessor
        )
        model.backbone.token_embedding.trainable = False
        model.backbone.position_embedding.trainable = False
    if args.fp8:
        model.quantize("float8")

    model.summary()
    optimizer, loss = get_optimizer_and_loss()
    model.compile(optimizer=optimizer, loss=loss, weighted_metrics=["accuracy"])
    callbacks = [
        GPUMemoryCallback(target_batches=[5, 10, 25, 50, 100, 150, 200, 300])
    ]
    model.fit(
        data, batch_size=args.batch_size, epochs=EPOCHS, callbacks=callbacks
    )
    if keras.backend.backend() == "tensorflow":
        model_memory_usage = callbacks[0].memory_usage
        print(f"GPU Memory Usage (in GB): {max(model_memory_usage)}")

    if args.fp8:
        from keras import layers

        count = 0
        for layer in model._flatten_layers(False, True):
            list_of_sublayers = list(layer._flatten_layers())
            if len(list_of_sublayers) == 1:  # leaves of the model
                if isinstance(layer, (layers.Dense, layers.EinsumDense)):
                    print(
                        layer.outputs_grad_scale.path,
                        layer.outputs_grad_scale.value,
                    )
                    count += 1
                if count > 10:
                    break

Issues

• The improvements in training speed are relatively small if the model is not sufficiently large. The recommendation is to train with a model containing parameters >5B.
• You will need hardware that supports fp8 tensor cores to gain the speedups.
• It is challenging to configure for Torch without low-level operations at this time. You can check the following links for float8 training in Torch instead:
  • https://github.com/pytorch-labs/float8_experimental
  • https://github.com/NVIDIA/TransformerEngine?tab=readme-ov-file#pytorch

References for Implementation

• https://github.com/NVIDIA/atex/tree/experimental-xla-fp8-layers/experimental/xla-fp8-layers
• https://github.com/google/flax/blob/main/flax/linen/fp8_ops.py

cc @kaixih

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 95.32164% with 16 lines in your changes are missing coverage. Please review.

Project coverage is 76.51%. Comparing base (5187fac) to head (197c6e7).
Report is 5 commits behind head on master.

Files	Patch %	Lines
keras/layers/core/dense.py	94.05%	2 Missing and 4 partials ⚠️
keras/layers/core/einsum_dense.py	94.54%	2 Missing and 4 partials ⚠️
keras/layers/core/embedding.py	92.00%	1 Missing and 1 partial ⚠️
keras/layers/layer.py	80.00%	0 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #19488      +/-   ##
==========================================
+ Coverage   76.27%   76.51%   +0.23%     
==========================================
  Files         367      367              
  Lines       41233    41578     +345     
  Branches     8076     8127      +51     
==========================================
+ Hits        31451    31813     +362     
+ Misses       8060     8052       -8     
+ Partials     1722     1713       -9     

Flag	Coverage Δ
keras	76.36% <94.15%> (+0.22%)	⬆️
keras-jax	60.54% <93.56%> (+0.24%)	⬆️
keras-numpy	54.36% <57.89%> (+0.08%)	⬆️
keras-tensorflow	61.77% <94.15%> (+0.22%)	⬆️
keras-torch	60.63% <92.39%> (+0.21%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[james77777778]

This PR needs to wait #19495

There might be potential bugs in LoRA and the quantization in torch backend if torch_params doesn't correctly track/untrack the added/removed variables.

[fchollet]

Thanks for the PR! 👍

[fchollet]

I think this should not be set on the layer as a constructor arg, but on the DTypePolicy. If you want to configure it, call the layer with layer.quantize(Float8DtypePolicy(amax_history_length=32)) or something like that. We may want to expose dedicate DTypePolicy subclasses to expose such arguments (this is the first one we encounter but there may be more of them in the future).

[james77777778]

I struggled to set args to DTypePolicy and I'm uncertain if this is a good approach.

In the current codebase, configuring DTypePolicy for individual layers is challenging (as shown in #19381 ).

1. DTypePolicy is treated as a global state. Making it difficult to set different amax_history_length for different layers. (We also need to consider scenarios involving subclasses)
2. The serialization/deserialization may be an issue.

Based on these concerns, I adopted the approach simiar to lora_rank that amax_history_length defaults to None and is overridden as 1024 in self._float8_build if not specified.

What do you think?

[fchollet]

I would strongly prefer to keep dtype configuration and layer constructor args separate. All dtype configuration should go through DTypePolicy.

We also need to consider scenarios involving subclasses

If the dtype policy is a settable attribute, this is not a big issue -- you can instantiate the subclass, and then set the dtype policy on a specific layer owned by the subclass.

For serialization, the dtype policy (and its args) should be part of the layer config, under the dtype entry in the config dict.

lora_rank is simply a shortcut to 1. instantiate layer, 2. call enable_lora(). It's not necessary from a configuration standpoint.

[james77777778]

Got it. The idea of keeping the configuration minimal for layers is solid.

As we discussed above, I have introduced a QuantizedFloat8DTypePolicy in this PR. I think it would clearer to use "Quantized" in the naming, as it involves quantizing and dequantizing.

BTW, I'd like to point it out that setting the dtype policy after the instantiation would not trigger the quantization process.
We need to either pass the correct dtype_policy (such as "float8_from_float32") though dtype argment in __init__ or utilize quantize api.

[fchollet]

BTW, I'd like to point it out that setting the dtype policy after the instantiation would not trigger the quantization process.

Do you think it would be a good idea to implement a setter that checks whether the policy is a quantized policy, and calls .quantize() if it is?

[james77777778]

Do you think it would be a good idea to implement a setter that checks whether the policy is a quantized policy, and calls .quantize() if it is?

This sounds like a great idea!
I think this change will also improve the UX of serializing and deserializing quantizable layers.

[james77777778]

I have added the setter of dtype_policy and it works well.

[james77777778]

After adding dtype_policy setter, we can now implement a cleaner solution to support the quantization of subclasses.

import tempfile

import numpy as np

import keras
from keras import dtype_policies


@keras.saving.register_keras_serializable("MyPackage")
class MySubclass(keras.Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.layer1 = keras.layers.Dense(4)
        self.layer2 = keras.layers.BatchNormalization(axis=-1)
        self.layer3 = keras.layers.ReLU()

    def call(self, inputs):
        x = self.layer1(inputs)
        x = self.layer2(x)
        return self.layer3(x)

    def get_config(self):
        config = super().get_config()
        quantized_dtype_policies = {}
        for idx, layer in enumerate(
            self._flatten_layers(include_self=False, recursive=False)
        ):
            if isinstance(
                layer.dtype_policy, dtype_policies.QuantizedDTypePolicy
            ):
                quantized_dtype_policies[str(idx)] = layer.dtype_policy
        if quantized_dtype_policies:
            config.update(
                {"quantized_dtype_policies": quantized_dtype_policies}
            )
        return config

    @classmethod
    def from_config(cls, config):
        quantized_dtype_policies = config.pop("quantized_dtype_policies", None)
        obj = super(MySubclass, cls).from_config(config)
        if quantized_dtype_policies:
            for idx, layer in enumerate(
                obj._flatten_layers(include_self=False, recursive=False)
            ):
                if str(idx) in quantized_dtype_policies:
                    layer.dtype_policy = quantized_dtype_policies[str(idx)]
        return obj


model = keras.Sequential([keras.layers.Input([8]), MySubclass()])
model.quantize("int8")

with tempfile.TemporaryDirectory() as tempdir:
    path = f"{tempdir}/model.keras"
    model.save(path)
    reloaded_model = keras.saving.load_model(path)

x = keras.random.uniform([1, 8])
y1 = model(x, training=False)
y2 = reloaded_model(x, training=False)
np.testing.assert_allclose(y1, y2)

[fchollet]

Excellent work -- the changes look good!