Mixed-precision training combines numerical formats during model training: it uses fast 16-bit computation where reduced precision is safe, while retaining 32-bit values for gradients, weights, and updates that need stability. This usually provides much of the speed benefit of float16 without meaningfully reducing model quality. section 18.3.3
Enable mixed precision in a mini-GPT
A Keras model can enable the policy before constructing its layers: config_set_dtype_policy("mixed_float16"). Most of the forward pass then computes in float16, while layer variables remain float32. An eight-layer decoder model can therefore use faster computation without storing and updating its weights in half precision. section 16.2.1 section 18.3.3
Apply the technique with these practical choices:
- Keep numerically unstable operations, such as softmax or cross-entropy, in float32 by passing dtype="float32" to the relevant layer. section 18.3.3
- Use loss scaling when small float16 gradients risk rounding to zero; multiplying the loss makes those gradients more representable. section 18.3.4
- Prefer automatic scaling when the appropriate factor may change during training: optimizer_adam(learning_rate = 1e-3) |> optimizer_loss_scale(). section 18.3.4
Listing 16.6 Creating a mini-GPT functional model
config_set_dtype_policy("mixed_float16") #1
vocab_size <- tokenizer$vocabulary_size()
hidden_dim <- 512
intermediate_dim <- 2056
num_heads <- 8
num_layers <- 8
inputs <- keras_input(shape = c(NA), dtype = "int32", name = "inputs")
embedding <-
layer_positional_embedding(, sequence_length, vocab_size, hidden_dim)
x <- inputs |>
embedding() |>
layer_layer_normalization()
for (i in seq_len(num_layers)) {
x <- x |>
layer_transformer_decoder(hidden_dim, intermediate_dim, num_heads)
}
outputs <- x |> embedding(reverse = TRUE)
mini_gpt <- keras_model(inputs, outputs)
Inspect a complete R/Keras mini-GPT configuration that enables mixed_float16 before defining the model. section 16.2.1
