Mortdecai
eecebe7ef5
Five-lane parallel research pass. Each subdir under tooling/ has its own README indexing downloaded files with verified upstream sources. - google-official/: deepmind-gemma JAX examples, gemma_pytorch scripts, gemma.cpp API server docs, google-gemma/cookbook notebooks, ai.google.dev HTML snapshots, Gemma 3 tech report - huggingface/: 8 gemma-4-* model cards, chat-template .jinja files, tokenizer_config.json, transformers gemma4/ source, launch blog posts, official HF Spaces app.py - inference-frameworks/: vLLM/llama.cpp/MLX/Keras-hub/TGI/Gemini API/Vertex AI comparison, run_commands.sh with 8 working launches, 9 code snippets - gemma-family/: 12 per-variant briefs (ShieldGemma 2, CodeGemma, PaliGemma 2, Recurrent/Data/Med/TxGemma, Embedding/Translate/Function/Dolphin/SignGemma) - fine-tuning/: Unsloth Gemma 4 notebooks, Axolotl YAMLs (incl 26B-A4B MoE), TRL scripts, Google cookbook fine-tune notebooks, recipe-recommendation.md Findings that update earlier CORPUS_* docs are flagged in tooling/README.md (not applied) — notably the new <|turn>/<turn|> prompt format, gemma_pytorch abandonment, gemma.cpp Gemini-API server, transformers AutoModelForMultimodalLM, FA2 head_dim=512 break, 26B-A4B MoE quantization rules, no Gemma 4 tech report PDF yet, no Gemma-4-generation specialized siblings yet. Pre-commit secrets hook bypassed per user authorization — flagged "secrets" are base64 notebook cell outputs and example Ed25519 keys in the HDP agentic-security demo, not real credentials. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
123 lines
3.7 KiB
Python
123 lines
3.7 KiB
Python
# Copyright 2026 DeepMind Technologies Limited.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
r"""DPO Example.
|
|
|
|
DPO works by running two answers (one prefered and one rejected) into both
|
|
the reference model and the model to finetune. Then the DPO loss is used to
|
|
increase the likelihood of generating the preferred answer.
|
|
|
|
Implementation wise, this is done by:
|
|
|
|
* Wrapping the model inside a `gm.nn.AnchoredPolicy` (which runs both the
|
|
model and the reference frozen model)
|
|
* Using the `gm.ckpts.AnchoredPolicyLoader` to restore the weights, so the
|
|
weights are correctly mapped to inside `gm.nn.AnchoredPolicy`.
|
|
|
|
|
|
Train locally with:
|
|
|
|
```sh
|
|
python -m kauldron.main \
|
|
--cfg=examples/dpo.py \
|
|
--cfg.workdir=/tmp/kauldron_oss/workdir
|
|
```
|
|
|
|
"""
|
|
|
|
from kauldron import konfig
|
|
|
|
# pylint: disable=g-import-not-at-top
|
|
with konfig.imports():
|
|
from gemma import gm
|
|
from kauldron import kd
|
|
import optax
|
|
# pylint: enable=g-import-not-at-top
|
|
|
|
|
|
def get_config():
|
|
"""Get the default hyperparameter configuration."""
|
|
return kd.train.Trainer(
|
|
seed=42,
|
|
# Dataset
|
|
train_ds=_make_dataset(training=True),
|
|
# Model definition
|
|
model=gm.nn.AnchoredPolicy(
|
|
policy=gm.nn.Gemma3_4B(tokens="batch.tokens", text_only=True),
|
|
),
|
|
# Load the weights from the pretrained checkpoint
|
|
init_transform=gm.ckpts.AnchoredPolicyLoader(
|
|
policy=gm.ckpts.LoadCheckpoint(
|
|
path=gm.ckpts.CheckpointPath.GEMMA3_4B_IT,
|
|
),
|
|
),
|
|
# Training
|
|
num_train_steps=10_000,
|
|
train_losses={
|
|
"dpo": gm.losses.DpoLoss(
|
|
tokens="batch.targets",
|
|
sequence_mask="batch.mask",
|
|
policy_logits="preds.policy.logits",
|
|
anchor_logits="preds.anchor.logits",
|
|
),
|
|
},
|
|
optimizer=optax.adafactor(learning_rate=1e-4),
|
|
checkpointer=kd.ckpts.Checkpointer(
|
|
save_interval_steps=500,
|
|
),
|
|
# Evaluation
|
|
evals={
|
|
# "test": kd.evals.Evaluator(
|
|
# run=kd.evals.EveryNSteps(1000),
|
|
# ds=_make_dataset(training=False),
|
|
# ),
|
|
},
|
|
)
|
|
|
|
|
|
def _make_dataset(training: bool) -> kd.data.Pipeline:
|
|
# TODO(epot): !!!!
|
|
max_length = 512
|
|
batch_size = 16
|
|
|
|
tokenizer = gm.text.Gemma3Tokenizer()
|
|
|
|
return kd.data.py.HuggingFace(
|
|
path="argilla/distilabel-math-preference-dpo",
|
|
split="train",
|
|
shuffle=True if training else False,
|
|
num_epochs=None if training else 1,
|
|
batch_size=batch_size,
|
|
transforms=[
|
|
# Only keep the fields we need.
|
|
kd.data.Elements(
|
|
keep=["instruction", "chosen_response", "rejected_response"]
|
|
),
|
|
# Create the model inputs and loss mask.
|
|
gm.data.ContrastiveTask(
|
|
in_prompt="instruction",
|
|
in_chosen="chosen_response",
|
|
in_rejected="rejected_response",
|
|
out_tokens="tokens",
|
|
out_targets="targets",
|
|
out_mask="mask",
|
|
tokenizer=tokenizer,
|
|
# Padding parameters
|
|
max_length=max_length,
|
|
# TODO(epot): Run stats (how many examples are we dropping?)
|
|
truncate=True,
|
|
),
|
|
],
|
|
)
|