No public description

MG_DOCKER_CODES_PIPER_ORIGIN_REV_ID: 955252640

.github/workflows/linter/requirements.txt

@@ -2,7 +2,7 @@ git+https://github.com/tensorflow/docs
 ipython
 jupyter
 nbconvert
-black==25.12.0
+black==26.3.1
 pyupgrade==3.21.2
 isort==7.0.0
 flake8==7.3.0

README.md

@@ -1,12 +1,12 @@
 # ![Google Cloud](https://avatars.githubusercontent.com/u/2810941?s=60&v=4) Google Cloud Vertex AI Samples
 
-This repository contains notebooks, code samples, sample apps, and other resources that demonstrate how to use, develop and manage machine learning and generative AI workflows using Google Cloud Vertex AI.
+This repository contains notebooks, code samples, sample apps, skills, and other resources that demonstrate how to use, develop and manage machine learning and generative AI workflows using Google Cloud Vertex AI.
 
 ## Overview
 
 [Vertex AI](https://cloud.google.com/vertex-ai) is a fully-managed, unified AI development platform for building and using generative AI. This repository is designed to help you get started with Vertex AI. Whether you're new to Vertex AI or an experienced ML practitioner, you'll find valuable resources here.
 
-For more Vertex AI Generative AI notebook samples, please visit the Vertex AI [Generative AI](https://github.com/GoogleCloudPlatform/generative-ai) GitHub repository.
+⚠️ For more Vertex AI Generative AI notebook samples, please visit the Vertex AI [Generative AI](https://github.com/GoogleCloudPlatform/generative-ai) GitHub repository.
 
 ## Explore, learn and contribute
 
@@ -16,11 +16,11 @@ You can explore, learn, and contribute to this repository to unleash the full po
 
 Explore this repository, follow the links in the header section of each of the notebooks to -
 
-![Colab](https://cloud.google.com/ml-engine/images/colab-logo-32px.png)  Open and run the notebook in [Colab](https://colab.google/)\
-![Colab Enterprise](https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png)  Open and run the notebook in [Colab Enterprise](https://cloud.google.com/colab/docs/introduction)\
-![Workbench](https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32)  Open and run the notebook in [Vertex AI Workbench](https://cloud.google.com/vertex-ai/docs/workbench/introduction)\
-![Github](https://cloud.google.com/ml-engine/images/github-logo-32px.png)  View the notebook on Github
-
+- Open and run the notebook in [Colab](https://colab.google/)
+- Open and run the notebook in [Colab Enterprise](https://cloud.google.com/colab/docs/introduction)
+- Open and run the notebook in [Vertex AI Workbench](https://cloud.google.com/vertex-ai/docs/workbench/introduction)
+- View the notebook on Github
+  
 ### Contribute
 
 See the [Contributing Guide](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/CONTRIBUTING.md).
@@ -35,7 +35,7 @@ To get started using Vertex AI, you must have a Google Cloud project.
 
 ## Repository structure
 
-```bash
+```text
 ├── notebooks
 │   ├── official - Notebooks demonstrating use of each Vertex AI service
 │   │   ├── automl
@@ -45,7 +45,23 @@ To get started using Vertex AI, you must have a Google Cloud project.
 │   │   ├── model_garden
 │   │   ├── ...
 ├── community-content - Sample code and tutorials contributed by the community
-
+├── docs - Deep-dive documentation and advanced setup guides
+└── skills - Suite of AI Agent "Skills" for Vertex AI
+    ├── README.md               # Developer guide for Vertex AI skills
+    ├── vertex-ai/              # Primary router for Vertex AI tasks
+    │   └── SKILL.md            # Entry point that routes across capabilities
+    ├── genai-sdk/              # Gemini API usage with Gen AI SDK
+    │   └── SKILL.md            # Guides for Python, JS/TS, Go, Java, C#
+    ├── vertex-deploy/          # Deploying models to Endpoints
+    │   └── SKILL.md            # Commands for open models & custom weights
+    ├── vertex-inference/       # Inferencing with GenAI models
+    │   └── SKILL.md            # Code samples for Gemini and OpenMaaS
+    └── vertex-tuning/          # Secondary router for model fine-tuning
+        ├── SKILL.md            # Router for tuning tasks
+        ├── gemini/             # Fine-tuning first-party Gemini models
+        │   └── SKILL.md
+        └── open-model/         # Fine-tuning third-party open models
+            └── SKILL.md
 ```
 ## Examples
 

community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/logger/requirements.txt

@@ -1,4 +1,4 @@
 google-cloud-bigquery==2.20.0
 tensorflow==2.12.1
-pillow==10.3.0
+pillow==12.2.0
 tf-agents==0.8.0

community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/prediction_container/requirements.txt

@@ -1,4 +1,4 @@
 google-cloud-pubsub==2.5.0
-pillow==10.3.0
+pillow==12.2.0
 tf-agents==0.8.0
 tensorflow==2.12.1

community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/simulator/requirements.txt

@@ -1,5 +1,5 @@
 dataclasses==0.6
 google-cloud-aiplatform==1.8.1
 tensorflow==2.12.1
-pillow==10.3.0
+pillow==12.1.1
 tf-agents==0.8.0

community-content/vertex-distributed-training/a3mega/llama-3-8b-nemo-pretraining/docker/requirements.txt

@@ -2,7 +2,7 @@ dllogger@git+https://github.com/NVIDIA/dllogger@v1.0.0
 
 # Fixing these libraries versions to avoid conflicting or broken packages.
 immutabledict==4.2.1
-protobuf==4.25.8
+protobuf==5.29.6
 opencv-python-headless==4.11.0.86
 docutils==0.16
 urllib3==2.6.3

community-content/vertex-distributed-training/a3mega/llama-3-8b-nemo-pretraining/requirements.txt

@@ -1,7 +1,7 @@
 absl-py==2.2.2
 annotated-types==0.7.0
 anyio==4.9.0
-black==25.1.0
+black==26.3.1
 cachetools==5.5.2
 certifi==2025.4.26
 charset-normalizer==3.4.2
@@ -31,14 +31,14 @@ packaging==25.0
 pathspec==0.12.1
 platformdirs==4.3.8
 proto-plus==1.26.1
-protobuf==5.29.4
-pyasn1==0.6.1
+protobuf==5.29.6
+pyasn1==0.6.3
 pyasn1_modules==0.4.2
 pydantic==2.11.4
 pydantic_core==2.33.2
 python-dateutil==2.9.0.post0
 pytz==2025.2
-requests==2.32.4
+requests==2.33.0
 rsa==4.9.1
 shapely==2.1.0
 six==1.17.0

docs/README.md

@@ -0,0 +1,10 @@
+# Agent Platform Training Clusters Blog Series 
+
+This directory contains deep-dive documentation, extended guides, and architectural references for Google Cloud Agent Platform Training Clusters.
+
+## Contents
+- **`vertex-training-cluster/`**: Documentation and setup guides for configuring and managing Agent Platform Training Clusters.
+
+## Blog Posts
+- [Model Distillation Best Practices](https://googlecloudplatform.github.io/vertex-ai-samples/vertex-training-cluster/model_distillation_best_practices): Explores off-policy model distillation, dataset curation, and hyperparameter scaling laws for training student models on Vertex AI.
+- [Forgetting Mitigation via Data Mixing](https://googlecloudplatform.github.io/vertex-ai-samples/vertex-training-cluster/forgetting_mitigation_data_mixing): Discusses catastrophic forgetting in model fine-tuning and how to mitigate it using multi-domain data mixing on Vertex AI.

docs/vertex-training-cluster/forgetting_mitigation_data_mixing.md

@@ -0,0 +1,448 @@
+<script type="text/javascript" async
+  src="https://cdn.mathjax.org/mathjax/latest/MathJax.js?config=TeX-MML-AM_CHTML">
+</script><br><br>
+
+# VTC Multi-Domain Dataset: Mitigating Catastrophic Forgetting with Data Mixing
+
+**Author:** [Mayank Sharan](mailto:mayanksharan@google.com)
+
+## Table of Contents
+
+* [Intro](#intro)
+* [Background](#background)
+* [Dataset Curation](#dataset-selection)
+* [Forgetting Mitigation Best Practices](#forgetting-mitigation-best-practices)
+    * [Experimental Setup](#experimental-setup)
+    * [Mitigating Forgetting](#mitigating-forgetting)
+    * [Mixing Ratios](#mixing-ratios)
+    * [Different Starting Models](#different-starting-models)
+* [Acknowledgements](#acknowledgements)
+* [References](#references)
+
+## Intro
+
+In this entry of our blog series on model training best practices for Vertex AI Training Cluster (VTC) customers, we talk about catastrophic forgetting and how to mitigate it. We focus on tuning public models using supervised fine tuning (SFT) with a specialized domain dataset. With both open and closed source models performing well on general tasks the primary goal of training one's own models is to improve the performance on specialized tasks. This typically comes at the cost of the model forgetting general capabilities which can severely limit the utility of the trained model.
+
+There are many possible interventions to limit forgetting, the most effective is mixing the target dataset with the actual dataset used in the model’s training. Since this is not available even for the most open source models, we have curated a multi-domain dataset that delivers the same benefits. This allows Vertex AI Training Cluster (VTC) customers to maintain and surpass frontier level model capabilities while training to further performance on specialized tasks.
+
+<figure align="center" id="fig-teaser">
+<table align="center" width="80%">
+  <tr>
+    <td align="center" width="100%">
+      <img src="images_data_mixing/teaser_forgetting.png" width="100%"><br>
+    </td>
+  </tr>
+</table>
+<figcaption align="left">
+<sub><b>Figure 1: Impact of mixing VTC Post Training dataset on Forgetting (8B model). </b> <i>Comparing SFT runs using only a specialized target dataset (MedMCQA) vs a mix of the target dataset and the VTC Post training dataset. Forgetting across all non-target domains is significantly mitigated with no performance loss on the target metric. Qwen3 Public here is the instruction tuned public Qwen3 8B model and the other two models are trained starting from the base Qwen3 8B model using only the target dataset and a mix of target dataset with the VTC dataset.</i></sub>
+</figcaption>
+</figure>
+
+We provide a thorough set of experiments to serve as a guide for reducing forgetting while post training the Qwen3 open-weight thinking model family, beginning from their base pre-trained checkpoints. Furthermore, we demonstrate the value our datasets provide across model sizes often surpassing the performance of the official Qwen3 models while preserving performance on the specialized task (See [Figure 1](#fig-teaser)). The Qwen3 family was specifically chosen for this study because its diverse range of parameter counts and the availability of both pre-trained and post-trained checkpoints provide an ideal environment for high-fidelity scaling analysis.
+
+To ensure our findings can be applied to a broad set of applications we validate our findings across five model sizes: 0.6B, 1.7B, 4B, 8B and 14B parameters. To support our VTC community in accelerating their own development, all code, datasets, and experiment configurations used in this blog are being made available for use in your training workloads.
+
+## Background
+
+Loss landscapes for neural networks have always been a complex multidimensional manifold rather than the simple convex ones that gradient descent is built for. Forgetting is a well known phenomenon in model customization, the first academically recorded instance being (McCloskey and Cohen, 1989) [<a href="#ref1">1</a>]. These manifolds have become even more complex with the introduction of Large Language Models where the number of parameters being optimized are typically in the billions. This makes it hard to mathematically grasp issues like forgetting. [Figure 2](#fig-loss-landscape) demonstrates a geometric understanding of why forgetting happens and how data mixing can mitigate it.
+
+<figure align="center" id="fig-loss-landscape">
+<table align="center" width="80%">
+  <tr>
+    <td align="center" width="100%">
+      <img src="images_data_mixing/background_loss_landscape.png" width="100%"><br>
+    </td>
+  </tr>
+</table>
+<figcaption align="left">
+<sub><b>Figure 2: Geometric Interpretation of Data Mixing to Mitigate Forgetting. </b> <i>Fine tuning objectives being meaningfully out of distribution from the pre-trained model often drives forgetting. Mixing in a dataset similar to the model distribution adjusts the objective enough to learn the new task without as much forgetting.</i></sub>
+</figcaption>
+</figure>
+
+## Dataset Selection
+
+Our primary requirements for a target dataset to run experiments to validate this were:
+
+1. It should be out of distribution to cause forgetting
+2. It should have an evaluation metric that it directly improves
+3. It should be able to train the model to perform better than the counterpart generalist model
+
+A good heuristic to determine where the data lies with respect to the model distribution is by calculating perplexity on samples from the dataset. Assuming
+- <span>$$X={x_1, x_2, \dots, x_N}$$</span> is a dataset sample represented as sequence of tokens
+- <span>$$P(x_i \mid x_{<i})$$</span> is the model likelihood of the i-th token given the sample till that token
+
+Then the perplexity for this sample can be calculated as follows:
+
+$$\begin{align*}  
+& ppl(X) = \exp \left( -\frac{1}{N} \sum_{i=1}^{N} \log P(x_i \mid x_{<i}) \right) \\
+& = \exp \left( -\frac{1}{N} \log (\prod_{i=1}^{N} P(x_i \mid x_{<i})) \right)
+\end{align*}  $$
+
+The product form of the equation shows that this is a direct measure of the joint probability of this sequence of tokens according to the model. Since this computation has a balancing negative sign to account for the negative log value a lower joint probability results in a higher perplexity value and vice versa. We evaluated the following datasets as out-of-distribution candidates:
+
+- [MedMCQA](https://huggingface.co/datasets/syz-ml2025/medmcqa) : Multiple Choice Questions (MCQ) dataset focusing on the medical domain
+- [BirdSQL](https://huggingface.co/datasets/birdsql/bird23-train-filtered) : Text to SQL generation dataset
+- [HardGen](https://huggingface.co/datasets/Bingguang/HardGen) : Function calling dataset
+
+We also calculate perplexity on [OpenR1-Math-220k](https://huggingface.co/datasets/open-r1/OpenR1-Math-220k) to provide a reference as we expect this to be in distribution for the model given the Qwen3 models are particularly strong in the math domain.
+
+<table id="tab-perplexity" style="margin-left:auto; margin-right:auto;">
+  <thead>
+    <tr>
+      <th>Dataset \ Model</th>
+      <th>Qwen3-0.6B</th>
+      <th>Qwen3-8B</th>
+      <th>Ours-0.6B</th>
+      <th>Ours-8B</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <td>MedMCQA</td>
+      <td>63.00</td>
+      <td>66.00</td>
+      <td>42.00</td>
+      <td>20.75</td>
+    </tr>
+    <tr>
+      <td>BirdSQL</td>
+      <td>38.50</td>
+      <td>55.50</td>
+      <td>45.50</td>
+      <td>17.00</td>
+    </tr>
+    <tr>
+      <td>HardGen</td>
+      <td>2.23</td>
+      <td>2.03</td>
+      <td>2.28</td>
+      <td>1.79</td>
+    </tr>
+    <tr>
+      <td>OpenR1-Math</td>
+      <td>8.63</td>
+      <td>9.75</td>
+      <td>6.44</td>
+      <td>5.34</td>
+    </tr>
+  </tbody>
+  <caption style="text-align: left;"><b>Table 1:</b> Perplexity score analysis with the public instruction tuned Qwen3 models and Qwen3 base models trained using the VTC dataset (Ours) to identify a suitable target dataset.</caption>
+</table>
+
+We see from [Table 1](#tab-perplexity) that OpenR1-Math-220k as we expected has low perplexity scores and HardGen shows an even lower perplexity score eliminating it from consideration. MedMCQA samples have high perplexity scores across all considered models. This dataset also has the advantage of a straightforward evaluation metric as we can use the validation split in the form of an MCQ verified evaluation.
+
+Based on this analysis we choose MedMCQA as our target dataset for these experiments. Additionally, since we are training a thinking model and the dataset does not have thinking traces we use the Qwen3-235B model to inject thinking traces into the training samples.
+
+## Forgetting Mitigation Best Practices
+
+### Experimental Setup
+
+#### Dataset Mixing
+
+We tested the impact of how forgetting responds to mixing the base dataset in different ratios with the target dataset. The base dataset here refers to the multi domain SFT dataset we have developed (see our [distillation blog post](https://googlecloudplatform.github.io/vertex-ai-samples/vertex-training-cluster/model_distillation_best_practices) [<a href="#ref2">2</a>] for details of the generation process) that can replicate and on certain metrics beat the public Qwen3 models. The target dataset here refers to the MedMCQA dataset. It is important to understand that in all mixing scenarios where the target dataset is present we will use the complete target dataset as that is the reasonable course of action we expect any customer to take. This leads to the total number of samples used in training varying based on the mixing ratio.
+
+We run 2 baseline experiments for each model size: using only the base dataset and only the target dataset. The mixing experiments are the base dataset being mixed in ratios of 0.9:0.1, 0.75:0.25 and 0.5:0.5. (0.9:0.1 means 90% of samples are from the base in-distribution dataset, while 10% are from the target out-of-distribution dataset.)
+
+The base dataset is randomly subsampled for each of these experiments. For simpler reference and analysis let’s define a mixing ratio <span>$$0 \le \alpha < 1$$</span>, such that the final dataset mixture includes <span>$$ N'_{B} = \frac{\alpha}{1 - \alpha} N_T$$</span> samples from the base dataset where <span>$$N_T$$</span> is the number of samples in the target dataset. In each of these mixtures the complete target dataset is used, contributing <span>$$N_T$$</span> samples for a total training dataset size of <span>$$\frac{N_T}{1 - \alpha}$$</span>.
+
+Since, our target dataset has 182,712 samples, this means that:
+
+- 0.9:0.1 ratio (<span>$$\alpha = 0.9$$</span>) : Uses a total of 1,827,120 training samples
+- 0.75:0.25 ratio (<span>$$\alpha = 0.75$$</span>) : Uses a total of 730,849 training samples
+- 0.5:0.5 ratio (<span>$$\alpha = 0.5$$</span>) : Uses a total of 365,425 training samples
+
+#### Evaluation
+
+<table id="tab-eval-setup" style="margin-left:auto; margin-right:auto;">
+  <thead>
+    <tr>
+      <th>Capabilities</th>
+      <th>Benchmarks</th>
+      <th># Test Samples</th>
+      <th>Eval Metrics</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <td rowspan="7">Math</td>
+      <td>AIME 24</td>
+      <td>30</td>
+      <td>pass@1 (average of 10)</td>
+    </tr>
+    <tr>
+      <td>AIME 25</td>
+      <td>30</td>
+      <td>pass@1 (average of 10)</td>
+    </tr>
+    <tr>
+      <td>BeyondAIME</td>
+      <td>100</td>
+      <td>pass@1 (average of 5)</td>
+    </tr>
+    <tr>
+      <td>Math 500</td>
+      <td>500</td>
+      <td>pass@1</td>
+    </tr>
+    <tr>
+      <td>HMMT 25</td>
+      <td>30</td>
+      <td>pass@1 (average of 10)</td>
+    </tr>
+    <tr>
+      <td>BRUMO 25</td>
+      <td>30</td>
+      <td>pass@1 (average of 10)</td>
+    </tr>
+    <tr>
+      <td>CMIMC 25</td>
+      <td>40</td>
+      <td>pass@1 (average of 10)</td>
+    </tr>
+    <tr>
+      <td rowspan="3">Science</td>
+      <td>GPQA</td>
+      <td>448</td>
+      <td>pass@1 (average of 5)</td>
+    </tr>
+    <tr>
+      <td>MMLU</td>
+      <td>14042</td>
+      <td>pass@1</td>
+    </tr>
+    <tr>
+      <td>MMLU Pro</td>
+      <td>12032</td>
+      <td>pass@1</td>
+    </tr>
+    <tr>
+      <td rowspan="2">Coding</td>
+      <td>HumanEval</td>
+      <td>164</td>
+      <td>pass@1 (average of 5)</td>
+    </tr>
+    <tr>
+      <td>LiveCodeBench v6</td>
+      <td>175</td>
+      <td>pass@1 (average of 5)</td>
+    </tr>
+    <tr>
+      <td>Instruction Following</td>
+      <td>IFEval</td>
+      <td>541</td>
+      <td>pass@1 (Strict Accuracy)</td>
+    </tr>
+    <tr>
+      <td>Reasoning</td>
+      <td>ARC-AGI 1</td>
+      <td>400</td>
+      <td>pass@1 (average of 5)</td>
+    </tr>
+    <tr>
+      <td>Medical (Target domain)</td>
+      <td>MedMCQA</td>
+      <td>4183</td>
+      <td>pass@1</td>
+    </tr>
+  </tbody>
+  <caption style="text-align: left;"><b>Table 2:</b> Comprehensive overview of task domains, evaluation benchmarks, and associated performance metrics.</caption>
+</table>
+
+Our evaluation benchmarks and metrics are detailed in [Table 2](#tab-eval-setup). To ensure statistical reliability on smaller datasets, we report metrics averaged over multiple independent runs to mitigate variance. For each domain with multiple evaluations, we utilize the average score across the core benchmarks as our primary performance indicator. To maintain a consistent comparison, both our trained models and the official Qwen3 thinking models were evaluated using standardized sampling parameters — `Temperature=0.6`, `Top-P=0.95`, `Top-K=20` and `Max-tokens=32768` — aligning with the recommended [best practices](https://huggingface.co/Qwen/Qwen3-14B#best-practices) from the official Qwen3 model card.
+
+Note that we have separated MedMCQA as a target metric instead of including it in the Science domain. This is to ensure clear outcomes from our experiments and to demonstrate impacts on model performance without any interference.
+
+#### Training
+
+##### Vertex AI Training Cluster
+
+All experiments and results presented were orchestrated using the [Vertex AI Training Cluster (VTC)](https://docs.cloud.google.com/vertex-ai/docs/training/training-clusters/overview). VTC is a managed Google Cloud service designed to simplify and accelerate large-scale AI workloads. It provides a simple managed user experience that enables optimized GPU scheduling, automated fault tolerance, high hardware resiliency, quick start recipes and science tooling which drastically reduces the time from cluster setup to production training and speeds up experimentation.
+
+##### Training Framework and Hyperparameters
+
+We utilize NVIDIA [NeMo RL](https://github.com/NVIDIA-NeMo/RL), an open library from the [NVIDIA NeMo framework](https://github.com/NVIDIA-NeMo/) as the primary training library, leveraging the Megatron backend for distributed scaling. Models are initialized from a Qwen3 Base checkpoint and fine-tuned with a 32,768 context window on curated datasets. Optimization is handled via AdamW (<span>$$\beta_1=0.9$$</span>, <span>$$\beta_2=0.95$$</span>, weight decay=0.1) using a linear warmup and cosine decay schedule. All training is conducted using BF16 mixed precision. There are many model sizes and dataset mixes used in the experimentation so the maximum learning rate is guided by learning rate scaling laws (see [distillation blog post](https://googlecloudplatform.github.io/vertex-ai-samples/vertex-training-cluster/model_distillation_best_practices#hyperparameter-scaling) [<a href="#ref2">2</a>] for more) available as a part of VTC. The value is validated by testing slight adjustments from the recommended value for each dataset mixture.
+
+### Mitigating Forgetting
+
+All models in this experiment are trained starting from the Qwen3 base checkpoint. We explore the impact of dataset mixing by comparing the public Qwen3 instruction tuned model performance with our two baselines — model trained with only the target dataset and model trained only with the base dataset — and with a model trained using a 0.9 ratio mix.
+
+<figure align="center" id="fig3_data_mixing">
+
+<table align="center" width="100%">
+  <tr>
+    <td align="center" width="50%">
+      <img src="images_data_mixing/fig3_math.png" width="100%"><br>
+      <sub><b>(a)</b> Math</sub>
+    </td>
+    <td align="center" width="50%">
+      <img src="images_data_mixing/fig3_science.png" width="100%"><br>
+      <sub><b>(b)</b> Science</sub>
+    </td>
+  </tr>
+  <tr>
+    <td align="center" width="50%">
+      <img src="images_data_mixing/fig3_coding.png" width="100%"><br>
+      <sub><b>(c)</b> Coding</sub>
+    </td>
+    <td align="center" width="50%">
+      <img src="images_data_mixing/fig3_ifeval.png" width="100%"><br>
+      <sub><b>(d)</b> IFEval</sub>
+    </td>
+  </tr>
+  <tr>
+    <td align="center" width="50%">
+      <img src="images_data_mixing/fig3_arc_agi.png" width="100%"><br>
+      <sub><b>(e)</b> ARC-AGI</sub>
+    </td>
+    <td align="center" width="50%">
+      <img src="images_data_mixing/fig3_medmcqa.png" width="100%"><br>
+      <sub><b>(f)</b> MedMCQA</sub>
+    </td>
+  </tr>
+</table>
+<figcaption align="left">
+<sub><b>Figure 3: Performance with and without Data Mixing.</b> <i>A comparison across (a) Math, (b) Science, (c) Coding, (d) IFEval, (e) ARC-AGI and (f) MedMCQA benchmarks showing how data mixing impacts forgetting and performance on the target metric.</i></sub>
+</figcaption>
+
+</figure>
+
+[Figure 3](#fig3_data_mixing) shows that for all non-target metrics other than Science using just the target dataset shows significant forgetting. Math and ARC-AGI are almost completely forgotten for all model sizes up to 8B parameters. The mixed dataset recovers the performance to similar levels as the base dataset. The base dataset delivers performance comparable to the public model in all domains and significantly better on ARC-AGI.
+
+The Science domain evaluations do not suffer severe forgetting likely because MedMCQA is very close to this domain. In fact, for the 8B and 14B sizes due to these transfer learning dynamics the <span>$$\alpha = 0.9$$</span> model outperforms both the public instruction-tuned and the base dataset (<span>$$\alpha = 1$$</span>) models.
+
+Performance on the target metric of MedMCQA follows expected behavior with best results achieved by the model when trained only with the target dataset. It is important to note that the <span>$$\alpha = 0.9$$</span> model for all sizes is still significantly better than the public instruction-tuned and base dataset (<span>$$\alpha = 1$$</span>) model and for all sizes other than the 0.6B mostly maintains the performance gains of the target dataset (<span>$$\alpha = 0$$</span>) model.
+
+#### Key Observations
+
+Combining these conclusions we can see that mixing with our base dataset:
+
+- Matches and outperforms the public instruction tuned model on general tasks.
+- Preserves the gains beyond the public model on target tasks.
+- Provides additional gains on tasks from a similar domain.
+
+### Mixing Ratios
+
+Now that we know that mixing the base dataset almost eliminates forgetting it is important to understand how performance changes for different mixing configurations. This is also important to examine as it determines training length and hence the cost. We will compare models trained only with the target dataset to models trained using dataset mixes with <span>$$\alpha = 0.5, 0.75, 0.9$$</span>. The ratio mentioned here refers to the proportion of the dataset from the base dataset.
+
+<figure align="center" id="fig4_mixing_ratios">
+
+<table align="center" width="100%">
+  <tr>
+    <td align="center" width="50%">
+      <img src="images_data_mixing/fig4_math.png" width="100%"><br>
+      <sub><b>(a)</b> Math</sub>
+    </td>
+    <td align="center" width="50%">
+      <img src="images_data_mixing/fig4_science.png" width="100%"><br>
+      <sub><b>(b)</b> Science</sub>
+    </td>
+  </tr>
+  <tr>
+    <td align="center" width="50%">
+      <img src="images_data_mixing/fig4_coding.png" width="100%"><br>
+      <sub><b>(c)</b> Coding</sub>
+    </td>
+    <td align="center" width="50%">
+      <img src="images_data_mixing/fig4_ifeval.png" width="100%"><br>
+      <sub><b>(d)</b> IFEval</sub>
+    </td>
+  </tr>
+  <tr>
+    <td align="center" width="50%">
+      <img src="images_data_mixing/fig4_arc_agi.png" width="100%"><br>
+      <sub><b>(e)</b> ARC-AGI</sub>
+    </td>
+    <td align="center" width="50%">
+      <img src="images_data_mixing/fig4_medmcqa.png" width="100%"><br>
+      <sub><b>(f)</b> MedMCQA</sub>
+    </td>
+  </tr>
+</table>
+<figcaption align="left">
+<sub><b>Figure 4: Performance across Mixing Ratios.</b> <i>A comparison across (a) Math, (b) Science, (c) Coding, (d) IFEval, (e) ARC-AGI and (f) MedMCQA benchmarks showing how dataset mixing ratios impact forgetting and performance on the target metric.</i></sub>
+</figcaption>
+
+</figure>
+
+[Figure 4](#fig4_mixing_ratios) shows that for all non-target metrics mixing helps achieve better performance than just using the target dataset even with a <span>$$\alpha = 0.5$$</span> mix. As expected the performance on non target metrics worsens as we lower the ratio of the base dataset. This effect is more pronounced in the smaller size models and for datasets like ARC-AGI where the mixed training provides a lot more gain. These patterns confirm that the gains on non target metrics are directly correlated to the base dataset.
+
+The effect while present for Science domain metrics is much less pronounced due to the cross domain characteristics. Even with lower ratios the performance for models 4B and larger holds, confirming that our target dataset of MedMCQA here contributes to limiting forgetting for this domain.
+
+The performance on the target metric, MedMCQA, stays mostly consistent with dips mostly when going from <span>$$\alpha = 0.75$$</span> mix to <span>$$\alpha = 0.5$$</span> mix. This aligns well as in all cases we are doing a complete epoch on the target dataset. The performance mostly holding at mixing ratios indicates that the tradeoff on the target metrics is relatively low even at an aggressive mixing ratio like 0.5.
+
+#### Key Observations
+
+The mixing ratio comparison shows us that:
+
+- A mixing ratio of 0.9 is the best for achieving gains on target tasks and limiting forgetting.
+- A mixing ratio of even 0.5 limits forgetting well while only doubling the token budget compared to training without any mixing.
+
+### Different Starting Models
+
+We have trained all our models starting from Qwen3 base checkpoints. A natural question here might be: What happens if we train starting from the instruction tuned public Qwen3 checkpoints for our target task? In this section we examine this question and compare the instruction-tuned model tuned with the target dataset and an <span>$$\alpha = 0.9$$</span> mix to the instruction-tuned model itself and the base model tuned with an <span>$$\alpha = 0.9$$</span> mix.
+
+<figure align="center" id="fig5_starting_models">
+
+<table align="center" width="100%">
+  <tr>
+    <td align="center" width="50%">
+      <img src="images_data_mixing/fig5_math.png" width="100%"><br>
+      <sub><b>(a)</b> Math</sub>
+    </td>
+    <td align="center" width="50%">
+      <img src="images_data_mixing/fig5_science.png" width="100%"><br>
+      <sub><b>(b)</b> Science</sub>
+    </td>
+  </tr>
+  <tr>
+    <td align="center" width="50%">
+      <img src="images_data_mixing/fig5_coding.png" width="100%"><br>
+      <sub><b>(c)</b> Coding</sub>
+    </td>
+    <td align="center" width="50%">
+      <img src="images_data_mixing/fig5_ifeval.png" width="100%"><br>
+      <sub><b>(d)</b> IFEval</sub>
+    </td>
+  </tr>
+  <tr>
+    <td align="center" width="50%">
+      <img src="images_data_mixing/fig5_arc_agi.png" width="100%"><br>
+      <sub><b>(e)</b> ARC-AGI</sub>
+    </td>
+    <td align="center" width="50%">
+      <img src="images_data_mixing/fig5_medmcqa.png" width="100%"><br>
+      <sub><b>(f)</b> MedMCQA</sub>
+    </td>
+  </tr>
+</table>
+<figcaption align="left">
+<sub><b>Figure 5: Performance across Starting Models.</b> <i>A comparison across (a) Math, (b) Science, (c) Coding, (d) IFEval, (e) ARC-AGI and (f) MedMCQA benchmarks showing how different starting models impact forgetting and performance on the target metric. Qwen 3 Public is the public instruction tuned Qwen3 model, &alpha;=0 (IT) and &alpha;=0.9 (IT) are the public instruction-tuned Qwen3 model trained only with the target dataset and the &alpha;=0.9 mixed dataset. &alpha;=0.9 (Base) is the base Qwen3 model trained on a 90% VTC dataset and 10% target dataset mix.</i></sub>
+</figcaption>
+
+</figure>
+
+In [Figure 5](#fig5_starting_models), among the non-target metrics other than science we see a common trend that starting with the IT model and using only the target dataset (<span>$$\alpha = 0$$</span>) shows severe forgetting. The base model and the instruction-tuned model trained using the <span>$$\alpha = 0.9$$</span> mix match or surpass the performance of the public model. This shows that starting with an instruction-tuned model while better than starting with the base model is still not a solution to forgetting. This also shows the high quality of our dataset that it can provide further gains on the public instruction-tuned model.
+
+Science domain metrics show different trends based on the model size. The advantage of data mixing is much more apparent in 0.6B and 1.7B models. Overall though there are no disadvantages to mixing across all model sizes. The IT model demonstrating significant forgetting is a clear indication that cross domain characteristics of our target dataset are not enough to mitigate forgetting on its own.
+
+The performance of the target metric, MedMCQA, shows no additional gain when we train using only the target dataset except for the 0.6B model, whether the starting model is a base model or the IT model. For all model sizes other than the 0.6B model we also see that the <span>$$\alpha = 0.9$$</span> mix trained model does not lose any meaningful performance compared to the target dataset only trained models. All the models trained using the target dataset clearly improve on the public model.
+
+#### Key Observations
+
+The comparison of different starting models shows us:
+
+- Using the instruction-tuned model as the starting model is better than the Base model.
+- The IT model also shows catastrophic forgetting and loses performance on non target metrics.
+- The <span>$$\alpha = 0.9$$</span> mix avoids forgetting even with the instruction-tuned starting model showing its robustness.
+
+## Acknowledgements
+
+We would like to express our sincere gratitude to the NVIDIA NeMo RL team–specifically Terry Kong– for their invaluable support throughout this project.
+
+We would also like to express our gratitude to our VTC teammates: Mohammadreza Mohseni, Weiran Zhao, Fei Xia, Youbao Tang, Xuehan Xiong, Joseph Pagadora, Jiuqiang Tang, Bo Wu, Lav Rai, and Minwoo Park for developing the underlying datasets, providing infrastructure support, feedback, and insightful discussions throughout the project. We also thank Ting Yu, Shengyang Dai, Peng Xu, and Saurabh Tiwary for their leadership and support.
+
+## References
+
+<a id="ref1"></a>[1] McCloskey, Michael, and Neal J. Cohen. "Catastrophic interference in connectionist networks: The sequential learning problem." Psychology of learning and motivation. Vol. 24. Academic Press, 1989. 109-165.
+
+<a id="ref2"></a>[2] Google Cloud. "Model Distillation Best Practices." Vertex AI Training Cluster Samples. Google, 2026. https://googlecloudplatform.github.io/vertex-ai-samples/vertex-training-cluster/model_distillation_best_practices.

docs/vertex-training-cluster/model_distillation_best_practices.md

@@ -0,0 +1,541 @@
+<script type="text/javascript" async
+  src="https://cdn.mathjax.org/mathjax/latest/MathJax.js?config=TeX-MML-AM_CHTML">
+</script><br><br>
+
+
+# Model Distillation Best Practices
+
+**Authors:** [Xuehan Xiong](mailto:xxman@google.com), [Youbao Tang](mailto:tangyoubao@google.com), [Fei Xia](mailto:feixia@google.com), [Bao Thach](mailto:baothach@google.com), [Joseph Pagadora](mailto:jcpagadora@google.com)
+
+
+## Table of Contents
+
+* [Intro](#intro)
+* [Background](#background)
+* [Dataset Curation](#dataset-curation)
+    * [Non-Agentic Tasks](#non-agentic-tasks)
+    * [Agentic Task (Tool Utilization)](#agentic-task-tool-utilization)
+* [Model Distillation Experiments](#model-distillation-experiments)
+    * [Experimental Setup](#experimental-setup)
+    * [Choosing Teacher Models](#choosing-teacher-models)
+    * [Number of Rollouts per Prompt](#number-of-rollouts-per-prompt)
+    * [Rejection Sampling](#rejection-sampling)
+* [Hyperparameter Scaling](#hyperparameter-scaling)
+    * [Scaling learning rates based on global batch size](#scaling-learning-rates-based-on-global-batch-size)
+    * [Scaling learning rates based on model parameters](#scaling-learning-rates-based-on-model-parameters)
+    * [Scaling learning rates based on token budget](#scaling-learning-rates-based-on-token-budget)
+* [Key Takeaways](#key-takeaways)
+* [Acknowledgements](#acknowledgements)
+* [Appendix](./appendix.md#appendix)
+
+## Intro
+
+Welcome to the inaugural installment of our blog series dedicated to model training best practices for Vertex AI Training Cluster (VTC) customers. In this article, we examine model distillation—a popular cost-effective methodology for optimizing student models by leveraging the intelligence of high-capacity teacher models. Two primary distillation schemes are typically employed: **on-policy** and **off-policy distillation**. In an on-policy setting, the student model generates its own reasoning traces during training, which are then evaluated or corrected by a teacher model in real-time. While effective, this approach is computationally intensive and requires constant active inference from the teacher.
+
+This blog focuses on **off-policy distillation**, a resource-efficient methodology where the student model is trained on a static, "gold-standard" dataset of reasoning traces previously curated by a teacher.
+While online distillation often requires complex orchestration—like offloading the student to CPU while the teacher scores the trajectories to save VRAM—the off-policy approach simplifies the workflow by **completely decoupling generation from training.** By leveraging frontier-level models like [Qwen3-235B](https://huggingface.co/Qwen/Qwen3-235B-A22B-Thinking-2507) or [GLM-4.7 355B](https://huggingface.co/zai-org/GLM-4.7-FP8) to generate high-quality trajectories upfront, developers can:
+
+* **Max out GPU Utilization:** Dedicate 100% of available VRAM and compute to the student's training phase without the overhead of model swapping.
+* **Scale Independently:** Generate datasets once and reuse them for multiple student architectures or hyperparameter sweeps.
+* **Simplify Orchestration:** Eliminate the need for multi-model memory management, allowing for a standard, high-throughput Supervised Fine-Tuning (SFT) pipeline.
+
+This allows developers to achieve "big model" reasoning logic in smaller, deployable students without the logistical headache of maintaining a live teacher-student link.
+
+<figure align="center" id="fig-teaser">
+
+<table align="center" width="90%">
+  <tr>
+    <td align="center" width="50%">
+      <img src="images/teaser_arc_agi.png" width="100%"><br>
+      <sub><b>(a)</b> ARC-AGI 1</sub>
+    </td>
+    <td align="center" width="50%">
+      <img src="images/teaser_tau2.png" width="100%"><br>
+      <sub><b>(b)</b> &tau;<sup>2</sup>-bench</sub>
+    </td>
+  </tr>
+</table>
+<figcaption align="left">
+<sub><b>Figure 1: Distilled Student Model Performance on Novel Domains.</b> <i>Comparison of student models fine-tuned via off-policy distillation versus official Qwen3 post-trained models of equivalent scale, showing significant performance gains on ARC-AGI 1 and &tau;<sup>2</sup>-bench.</i></sub>
+</figcaption>
+
+</figure>
+
+We provide a rigorous, step-by-step framework for reproducing the advanced reasoning capabilities of the Qwen3 open-weight thinking model family, beginning from their base pre-trained checkpoints. Furthermore, we demonstrate how this same distillation pipeline can be applied to novel domains, such as [ARC-AGI 1](https://arcprize.org/arc-agi/1/) and [<span>$$\tau^2$$</span>-bench](https://github.com/sierra-research/tau2-bench), to develop student models that surpass the performance of official Qwen3 variants of equivalent scale (See [Figure 1](#fig-teaser)). The Qwen3 family was specifically chosen for this study because its diverse range of parameter counts and the availability of both pre-trained and post-trained checkpoints provide an ideal environment for high-fidelity scaling analysis.
+
+To ensure the broad applicability of our findings, we conducted rigorous evaluations across four distinct task domains: competitive mathematics, instruction following, complex puzzle-solving, and tool utilization. We further validated these results across four model scales—1.7B, 4B, 8B, and 14B parameters—to demonstrate that our methodology remains consistent as model complexity increases.
+
+To support our VTC community in accelerating their own development, all code, datasets, and experiment configurations used in this blog are being made available for use in your training workloads.
+
+## Background
+
+To establish a mathematical foundation for our experiments, we first delineate the key differences between the two predominant strategies in model distillation: off-policy distillation and on-policy distillation.
+
+In off-policy distillation for an LLM, we assume:
+
+* <span>$$z \sim P(Z)$$</span>: prompts drawn from some prompt distribution
+* <span>$$x \sim P(X|z)$$</span>: responses generated by the teacher distribution conditioned on prompt <span>$$z$$</span>
+* <span>$$Q(X|z)$$</span>: student distribution we want to train
+
+The standard off-policy distillation objective minimizes the KL divergence from the student to the teacher:
+
+$$\begin{align*}  
+& \min_Q \mathbb{E}_{Z} \left[\mathbb{E}_{X|Z}\left[\log \frac{P(X|Z)}{Q(X|Z)}\right]\right] \\  
+&= \min_Q \sum_{z} P(z) \sum_{x}P(x|z) \log\frac{P(x|z)}{Q(x|z)}  
+\end{align*}  $$
+
+Since the teacher distribution (<span>$$P$$</span>) is fixed, minimizing KL is equivalent to:
+
+$$\begin{align*}  
+& \max_Q \mathbb{E}_{Z}\left[\mathbb{E}_{X|Z}\left[\log Q(X|Z)\right]\right] \\  
+&= \max_Q \sum_{z} P(z) \sum_{x}P(x|z) \log Q(x|z)  
+\end{align*}  $$
+
+Given a dataset of prompts and teacher-generated responses:
+
+$$\begin{align*}  
+\{z_i\}_{i=1}^M \quad & \text{where} \quad z_i \sim P(Z) \\  
+\{x_{ij}\}_{i,j=1}^{M,N} \quad & \text{where} \quad x_{ij} \sim P(X | z_i)  
+\end{align*}  $$
+
+the empirical loss by Monte Carlo sampling becomes:
+
+$$L_{\text{distill}}(Q) \approx -\frac{1}{MN}\sum_{i=1}^M\sum_{j=1}^N \log Q(x_{ij} | z_i).  $$
+
+This is simply maximum likelihood estimation for <span>$$Q$$</span> on teacher responses, which shares the same objective as Supervised Fine-tuning (SFT).
+
+In on-policy distillation, responses are sampled from the student:
+
+$$z \sim P(Z), \quad x \sim Q(X | z)  $$
+
+The teacher is only used to evaluate those student samples, so expectations are taken under <span>$$Q$$</span>, not <span>$$P$$</span>.
+
+The natural objective is:
+
+$$\min_Q \mathbb{E}_{Z} \left[\mathrm{KL}\big(Q(X|Z)||P(X|Z)\big) \right]  $$
+
+This objective defines the **reverse KL divergence**, which is characterized by its mode-seeking behavior. In this regime, the student model tends to concentrate its probability mass on the primary modes of the teacher distribution. This stands in contrast to the forward KL divergence used in off-policy distillation, which exhibits mean-seeking or mass-covering behavior, forcing the student to cover the entire support of the teacher’s distribution. Forward KL forces the student to allocate probability mass to *all* teacher modes, even those it cannot represent well. Under capacity constraints, this mass-covering approach produces a compromise distribution that can underperform a smaller, sharper target. This provides the intuition for our empirical results on Capacity Matching (Section [Choosing Teacher Models](#choosing-teacher-models)), where a same-sized teacher model proved most effective in some tasks.
+
+## Dataset Curation
+
+To facilitate the distillation of frontier-level reasoning, we established a high-fidelity data curation pipeline tailored to our four primary task domains: competitive mathematics, instruction following, reasoning/pattern recognition, and agentic tool utilization.
+
+### Non-Agentic Tasks
+
+#### Math
+
+We selected [OpenR1-Math](https://huggingface.co/datasets/open-r1/OpenR1-Math-220k) (default subset) as our primary prompt source and implemented a multi-stage filtering pipeline to ensure the highest data fidelity:
+
+1. **Verification filtering**: To ensure objective evaluation, we retained only "math-word-problem" types, discarding Multiple Choice Questions (MCQ) and proofs. MCQs were specifically excluded to mitigate the risk of the model arriving at a correct answer (25% baseline probability) through flawed reasoning chains.
+2. **Near-duplicate removal**: We employed Locality-Sensitive Hashing (LSH) to identify and prune near-identical prompts within the training set and across our evaluation benchmarks, preventing data contamination and overfitting.
+3. **Instructional sanitization**: We identified and removed hundreds of prompts containing extraneous translation instructions. This step ensures the student model remains focused on the mathematical reasoning task rather than defaulting to secondary objectives.
+4. **Solution leakage prevention**: To enforce authentic problem-solving, we stripped prompts containing pre-existing solutions, which would otherwise provide the teacher model with an "open-book" advantage and degrade the quality of the distilled reasoning traces.
+
+This rigorous curation process successfully refined the initial pool of 93,733 candidates into a high-quality dataset of 75,726 prompts.
+
+#### Instruction Following
+
+We selected the default partition of the [ifeval-like-data](https://huggingface.co/datasets/argilla/ifeval-like-data) dataset, comprising 550,000 unfiltered synthetic rows. To ensure data integrity, we applied a multi-stage refinement pipeline:
+
+1. **Invalid sample pruning**: We discarded rows with missing language codes or malformed JSON within the "kwargs" field to maintain structural consistency.
+2. **Conflict resolution**: We identified and removed pairs of mutually exclusive instructions that cannot be reliably evaluated together, utilizing a predefined mapping of instruction conflicts (`IFEVAL_INSTRUCTION_CONFLICTS`).
+3. **Adherence verification**: Each teacher-generated response was rigorously assessed using the [lm_eval](https://github.com/EleutherAI/lm-evaluation-harness) library. Prompts that failed to meet their defined constraints were excluded.
+4. **Strict accuracy filtering**: As a final quality gate, we retained only those samples where the response achieved "strict accuracy" at the prompt level, ensuring the student model learns from perfect examples of instruction following.
+5. **LSH-based deduplication**: We utilized LSH to prune near-duplicate prompts within the training set and across the official IFEval benchmark to prevent contamination.
+
+This pipeline successfully distilled the initial pool into 70,373 high-fidelity samples for our instruction-following training set.
+
+#### Reasoning/Pattern Recognition
+
+The [re-arc](https://github.com/michaelhodel/re-arc?tab=readme-ov-file) repository provides a way to programmatically synthesize [ARC-AGI-1](https://arcprize.org/arc-agi/1/) data. For each of the 400 training examples in the official ARC-AGI-1 dataset, re-arc provides a generator function to create similar puzzles following the same pattern (See [Figure 2](#fig-rearc-example) for one example). In total, we have generated 7926 puzzles for our experiments where we reserve 256 samples for validation and the rest for training.
+
+<figure align="center" id="fig-rearc-example">
+
+<table align="center" width="80%">
+  <tr>
+    <td align="center" width="100%">
+      <img src="images/arc_agi_original.png" width="100%"><br>
+      <sub><b>(a)</b> ARC-AGI original puzzles</sub>
+    </td>
+  </tr>
+  <tr>
+    <td align="center" width="510%">
+      <img src="images/arc_agi_generated.png" width="100%"><br>
+      <sub><b>(b)</b> Generated puzzles using re-arc</sub>
+    </td>
+  </tr>
+</table>
+<figcaption align="left">
+<sub><b>Figure 2:</b> <i>Example of ARC-AGI 1 puzzle synthesis using the re-arc repository, showing an original training example and a generated similar puzzle following the same pattern.</i></sub>
+</figcaption>
+
+</figure>
+
+#### Response Generation
+
+Following prompt collection, we utilize a "thinking" teacher model to generate reasoning traces. To maintain a lean data pipeline, we store only the sampled tokens; given our ~150K vocabulary size, persisting full logits or log-probabilities would create prohibitive storage overhead. This approach is statistically grounded: as the number of samples increases, it provides an unbiased estimate of the KL divergence from the student model to the teacher model.
+
+To maximize response diversity, we set both `Temperature` and `Top-P` to 1.0 during sampling. Finally, we prune any responses truncated by the maximum sequence length, as these instances frequently exhibit repetitive patterns that could degrade the student model’s performance.
+
+### Agentic Task (Tool Utilization)
+
+We utilized a specialized two-stage generation framework to synthesize high-complexity tool-use data for distillation training, leveraging sandboxed execution environments.
+
+1. **Task Generation**: This initial phase analyzes a target agent's specific tool list to propose a variety of diverse, high-level topics. For each identified topic, the system synthesizes a comprehensive user scenario that includes the initial environment status, the necessary database state, and precise evaluation criteria required for verification.
+2. **Trajectory Generation and Verification**: A verified <span>$$\tau^2$$</span>-bench sandbox is employed to execute each generated task several times in parallel, capturing a wide variety of trajectories. These execution outputs—comprising model responses, tool invocations, and subsequent state modifications—undergo a rigorous verification process. By applying deterministic checks such as action matching, database state differentials, and natural language assertions, the pipeline calculates the reward for every trajectory produced.
+
+Once trajectories are verified, they are carefully remapped into final training configurations to maximize learning efficiency. For distillation, the complete reasoning trace is captured and enclosed within required `<think>` tags, ensuring architectural consistency for the thinking model. A sample task and trajectory are provided in the [Appendix: <span>$$\tau^2$$</span>-bench Synthetic Example](./appendix.md#-bench-synthetic-example).
+
+## Model Distillation Experiments
+
+### Experimental Setup
+
+#### Evaluation
+Our evaluation benchmarks and metrics are detailed in [Table 1](#tab-eval-setup) and the prompts can be found in the [Appendix: Prompts Used in Evaluation](./appendix.md#prompts-used-in-evaluation). To ensure statistical reliability on smaller datasets, we report metrics averaged over multiple independent runs to mitigate variance. For the Mathematics domain, we utilize the average score across six core benchmarks as our primary performance indicator, while granular results for individual benchmarks are provided in the [Appendix: Individual Math Benchmark Results](./appendix.md#individual-math-benchmark-results). For <span>$$\tau^2$$</span>-bench, we use GLM-4.7-FP8 as the user LLM and report the average score across three domains, “Telecom”, “Retail”, and “Airline”. To maintain a consistent comparison, both our distilled student models and the official Qwen3 thinking models were evaluated using standardized sampling parameters—`Temperature=0.6`, `Top-P=0.95`, and `Top-K=20`—aligning with the recommended [best practice](https://huggingface.co/Qwen/Qwen3-14B#best-practices) from the official Qwen3 model card.
+
+<table id="tab-eval-setup" style="margin-left:auto; margin-right:auto;">
+  <thead>
+    <tr>
+      <th>Capabilities</th>
+      <th>Benchmarks</th>
+      <th># Test Samples</th>
+      <th>Eval Metrics</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <td rowspan="6">Math</td>
+      <td>AIME 24</td>
+      <td>30</td>
+      <td>pass@1 (average of 16)</td>
+    </tr>
+    <tr>
+      <td>AIME 25</td>
+      <td>30</td>
+      <td>pass@1 (average of 16)</td>
+    </tr>
+    <tr>
+      <td>BeyondAIME</td>
+      <td>100</td>
+      <td>pass@1 (average of 5)</td>
+    </tr>
+    <tr>
+      <td>HMMT 25</td>
+      <td>30</td>
+      <td>pass@1 (average of 16)</td>
+    </tr>
+    <tr>
+      <td>BRUMO 25</td>
+      <td>30</td>
+      <td>pass@1 (average of 16)</td>
+    </tr>
+    <tr>
+      <td>CMIMC 25</td>
+      <td>40</td>
+      <td>pass@1 (average of 16)</td>
+    </tr>
+    <tr>
+      <td>Instruction Following</td>
+      <td>IFEval</td>
+      <td>541</td>
+      <td>pass@1 (Strict Accuracy)</td>
+    </tr>
+    <tr>
+      <td>Reasoning</td>
+      <td>ARC-AGI 1</td>
+      <td>400</td>
+      <td>pass@1 (average of 5)</td>
+    </tr>
+    <tr>
+      <td>Tool use</td>
+      <td>&tau;<sup>2</sup>-bench</td>
+      <td>278</td>
+      <td>pass@1 (average of 4)</td>
+    </tr>
+  </tbody>
+  <caption style="text-align: left;"><b>Table 1:</b> Comprehensive overview of task domains, evaluation benchmarks, and associated performance metrics.</caption>
+</table>
+
+#### Training
+
+**Vertex AI Training Cluster**
+
+To orchestrate the computational demands of our experiments, we utilized the [Vertex AI Training Cluster (VTC)](https://docs.cloud.google.com/vertex-ai/docs/training/training-clusters/overview). VTC is a managed Google Cloud service designed to simplify and accelerate large-scale AI workloads. It provides a familiar, open-source Slurm user experience that enables optimized GPU scheduling, automated fault tolerance, and high hardware resiliency, which drastically reduces the time from cluster setup to production training.
+
+Our training infrastructure leverages VTC's high-performance [compute resources](https://docs.cloud.google.com/vertex-ai/docs/training/training-clusters/compute-resources), specifically utilizing A3-Mega (NVIDIA H100 GPUs), A3-Ultra (NVIDIA H200 GPUs), and A4 GPU (NVIDIA HGX B200) platforms powered by NVIDIA. To handle the communication overhead of distributed training, node connectivity is highly optimized for each hardware generation:
+
+* A3-Mega clusters utilize [GPUDirect-TCPXO](https://docs.cloud.google.com/compute/docs/gpus/gpudirect) for low-latency host-bypass networking
+* A3-Ultra and A4 clusters leverage [RoCE v2 (RDMA over Converged Ethernet)](https://cloud.google.com/blog/products/networking/rdma-rocev2-for-ai-workloads-on-google-cloud).
+
+By leveraging network topologies specifically optimized for training on large clusters of GPUs, this environment provides the high throughput and scaling efficiency necessary to reliably train and finetune frontier-level models.
+
+**Training Framework and Hyperparameters**
+
+We utilize NVIDIA [NeMo RL](https://github.com/NVIDIA-NeMo/RL) as the primary training framework, leveraging a Megatron backend for distributed scaling. We implement the <span>$$\tau^2$$</span>-bench sandbox environment inside NVIDIA [NeMo Gym](https://github.com/NVIDIA-NeMo/Gym), which provides a unified interface for building and scaling reinforcement learning environments and seamlessly integrated with the NeMo RL library for RL training runs.
+
+Models are initialized from a Qwen3 Base checkpoint and fine-tuned with a 32,768 context window on curated datasets. Optimization is handled via AdamW (<span>$$\beta_1=0.9$$</span>, <span>$$\beta_2=0.95$$</span>, weight decay=0.1) using a linear warmup and cosine decay schedule. To manage computational load, we employ tensor parallelism (2-way for 1.7B–8B models; 4-way for 14B) alongside sequence parallelism, activation checkpointing, and ZeRO-2. For further reading on parallelization strategies, see this [ultrascale playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook). All training is conducted using BF16 mixed precision.
+
+### Choosing Teacher Models
+
+A critical decision in the distillation pipeline is the selection of an appropriate teacher model for a given task domain. While conventional wisdom often suggests that "bigger is better," our empirical results across four benchmarks (illustrated in [Figure 3](#fig-teacher-student-matrix)) reveal a more nuanced landscape. Specifically, on well-defined reasoning tasks like Mathematics and IFEval, capacity-matched (same-sized) teachers frequently outperform their larger counterparts. Conversely, on novel or highly complex domains like ARC-AGI and <span>$$\tau^2$$</span>-bench, massive teacher models remain the superior choice. Below, we provide a formal derivation to explain this capacity-matching phenomenon and the trade-offs between approximation bias and teacher error.
+
+<figure align="center" id="fig-teacher-student-matrix">
+
+<table align="center" width="100%">
+  <tr>
+    <td align="center" width="50%">
+      <img src="images/teacher_student_matrix_ifeval.png" width="100%"><br>
+      <sub><b>(a)</b> IFEval</sub>
+    </td>
+    <td align="center" width="50%">
+      <img src="images/teacher_student_matrix_math.png" width="100%"><br>
+      <sub><b>(b)</b> Math Average</sub>
+    </td>
+  </tr>
+  <tr>
+    <td align="center" width="50%">
+      <img src="images/teacher_student_matrix_arc_agi1.png" width="100%"><br>
+      <sub><b>(c)</b> ARC-AGI 1</sub>
+    </td>
+    <td align="center" width="50%">
+      <img src="images/teacher_student_matrix_tau2.png" width="100%"><br>
+      <sub><b>(d)</b> &tau;<sup>2</sup>-bench</sub>
+    </td>
+  </tr>
+</table>
+<figcaption align="left">
+<sub><b>Figure 3: Teacher-Student Distillation Performance Matrices.</b> <i>A comparison of distillation outcomes across (a) IFEval, (b) Math, (c) ARC-AGI 1, and (d) &tau;<sup>2</sup>-bench benchmarks. The color intensity represents accuracy (%).</i></sub>
+</figcaption>
+
+</figure>
+
+Let <span>$$Q$$</span> and <span>$$P$$</span> denote the student and teacher distribution, and the student model class (e.g., all 8B models following the same architecture) is <span>$$\mathbf{Q}_{8B}$$</span>.
+
+**Optimal teacher model in theory**
+
+Suppose the student is an 8B model. The best teacher model is the optimal 8B model <span>$$Q^*$$</span> where <span>$$Q^* \in \mathbf{Q}_{8B}$$</span> for this task because if we have infinite data and perfect optimization, distillation can recover this optimal model exactly.
+
+**Why a larger teacher can still help**
+
+Now suppose the teacher is a larger model (<span>$$P_L$$</span>). The student solves
+
+$$Q^* = \text{argmin}_{Q\in\mathbf{Q}_{8B}} \mathrm{KL}(P_L|Q)$$
+
+This is the best 8B approximation of the larger teacher. Two competing effects appear:
+
+1. **Approximation bias**
+
+   Because the student is capacity-limited,
+
+   $$\inf_{Q \in \mathbf{Q}_{8B}} \mathrm{KL}(P_L|Q) > 0$$
+
+   So distillation from a very rich teacher may force the student to approximate a distribution it cannot represent well. This is the "mean-seeking / mass-covering" mentioned in the [Background](#background) section.
+
+2. **Teacher suboptimality**
+
+   In practice we rarely have <span>$$Q^*$$</span>, the true optimal 8B model. Instead we have a trained 8B model <span>$$\hat{Q}$$</span>, which contains optimization error and data error. A larger teacher (<span>$$P_L$$</span>) may actually be closer to the true distribution (<span>$$P^*$$</span>).
+   If
+
+   $$\mathrm{KL}(P^*|P_L) < \mathrm{KL}(P^*|\hat{Q})$$
+
+   then projecting (<span>$$P_L$$</span>) onto the 8B class can produce a better 8B model than the original 8B model.
+
+When an 8B student uses an 8B teacher, the projection error is inherently small due to matched capacity. However, if that 8B teacher is poorly optimized (e.g., on benchmarks like ARC-AGI and <span>$$\tau^2$$</span>-bench), its high **Teacher Error** dominates. Conversely, a massive model like Qwen3-235B or GLM 4.7, even if it has a higher **Projection Error** due to the size difference, can significantly lower the **Teacher Error** because it holds a more accurate approximation of the true distribution.
+
+### Number of Rollouts per Prompt
+
+We employ Monte Carlo sampling to approximate KL divergence, where expanding either the prompt set or the number of rollouts per prompt serves to reduce the variance of the estimate. However, high-fidelity prompts are often a finite resource—particularly in the Mathematics domain, which is constrained by the historical volume of competitive math problems. To compensate, we increase the number of teacher rollouts per prompt. This approach captures the teacher model’s **inherent uncertainty and multi-modal behavior** (e.g., discovering multiple valid reasoning paths to the same solution), enabling the student to map the full probability landscape rather than converging on a single, isolated trajectory.
+
+<figure align="center" id="performance_vs_rollouts">
+
+<table align="center" width="100%">
+  <tr>
+    <td align="center" width="50%">
+      <img src="images/response_scaling_ifeval.png" width="100%"><br>
+      <sub><b>(a)</b> IFEval</sub>
+    </td>
+    <td align="center" width="50%">
+      <img src="images/response_scaling_math.png" width="100%"><br>
+      <sub><b>(b)</b> Math Average</sub>
+    </td>
+  </tr>
+  <tr>
+    <td align="center" width="50%">
+      <img src="images/response_scaling_arc_agi.png" width="100%"><br>
+      <sub><b>(c)</b> ARC-AGI 1</sub>
+    </td>
+    <td align="center" width="50%">
+      <img src="images/response_scaling_tau2.png" width="100%"><br>
+      <sub><b>(d)</b> &tau;<sup>2</sup>-bench</sub>
+    </td>
+  </tr>
+</table>
+<figcaption align="left">
+<sub><b>Figure 4: Performance vs. Number of Rollouts.</b> <i>A comparison across (a) IFEval, (b) Math, (c) ARC-AGI, and (d) &tau;<sup>2</sup>-bench benchmarks showing how performance scales as the number of teacher rollouts per prompt increases from 1 to 16.</i></sub>
+</figcaption>
+
+</figure>
+
+[Figure 4](#performance_vs_rollouts) demonstrates that performance on the Math, ARC-AGI 1, and &tau;<sup>2</sup>-bench domains monotonically increases as the number of rollouts increases, while IFEval performance saturates at 8 rollouts.
+
+### Rejection Sampling
+
+This section explores whether rejection sampling on teacher responses—specifically, pruning trajectories that yield incorrect answers—enhances distillation performance. Formally, this approach **minimizes the KL divergence** against a **reweighted teacher distribution**, where incorrect paths are zero-weighted and valid paths are renormalized. [Figure 5](#fig_rejection_sampling) evaluates three strategies: utilizing the full response set, randomly subsampling to match the count of correct responses, and isolating correct responses only. Detailed acceptance rates for these task-teacher pairings are cataloged in [Table 2](./appendix.md#acceptance-rates-for-rejection-sampling) within the Appendix.
+
+<figure align="center" id="fig_rejection_sampling">
+
+<table align="center" width="100%">
+  <tr>
+    <td align="center" width="50%">
+      <img src="images/rejection_sampling_comparison_ifeval.png" width="100%"><br>
+      <sub><b>(a)</b> IFEval</sub>
+    </td>
+    <td align="center" width="50%">
+      <img src="images/rejection_sampling_comparison_math.png" width="100%"><br>
+      <sub><b>(b)</b> Math Average</sub>
+    </td>
+  </tr>
+  <tr>
+    <td align="center" width="50%">
+      <img src="images/rejection_sampling_comparison_arcagi.png" width="100%"><br>
+      <sub><b>(c)</b> ARC-AGI 1</sub>
+    </td>
+    <td align="center" width="50%">
+      <img src="images/rejection_sampling_comparison_tau2.png" width="100%"><br>
+      <sub><b>(d)</b> &tau;<sup>2</sup>-bench</sub>
+    </td>
+  </tr>
+</table>
+<figcaption align="left">
+<sub><b>Figure 5: Performance with and without Rejection Sampling.</b> <i>A comparison across (a) IFEval, (b) Math, (c) ARC-AGI, and (d) &tau;<sup>2</sup>-bench benchmarks showing how distillation performance is impacted by the use of rejection sampling during data curation.</i></sub>
+</figcaption>
+
+</figure>
+
+For instruction following (IFEval), rejection sampling provides a significant performance uplift within the same-sample-count regime. Notably, at the 14B model scale, this technique improves the student model by 2 percentage points compared to training on the full response set, despite utilizing a smaller volume of data. Conversely, for the Mathematics, ARC-AGI 1, <span>$$\tau^2$$</span>-bench domains, rejection sampling does not yield performance gains, with the highest accuracy achieved by utilizing all available teacher responses. This suggests that while rejection sampling refines the training distribution, it may also inadvertently prune "near-miss" cases or highly challenging problems that are essential for developing robust reasoning capabilities in those specific domains. For synthetic datasets like <span>$$\tau^2$$</span>-bench, imperfections in automated evaluation criteria may also cause the rejection of trajectories that contain high-quality reasoning traces despite an incorrect final answer.
+
+## Hyperparameter Scaling
+
+In practice, development-phase experimentation rarely mirrors the scale of final model training. To accelerate iteration, developers often conduct ablation studies using reduced token budgets and smaller model architectures. To help bridge this gap, we have compiled several **key rules of thumb** for translating hyperparameters from these 'proxy' settings to your final, full-scale production runs. All scaling studies below are conducted using the [IFEval-like](#instruction-following) dataset.
+
+### Scaling learning rates based on global batch size
+
+To accelerate training throughput, the most direct lever is scaling GPU resources and increasing the **Global Batch Size (B)**. However, effective scaling requires more than just hardware; the learning rate (<span>$$\eta$$</span>) must be precisely adjusted in tandem with the batch size to maintain an optimal convergence trajectory.
+
+
+<figure align="center" id="fig-lr-scale-gbs">
+<table align="center" width="80%">
+  <tr>
+    <td align="center" width="100%">
+      <img src="images/bs_scaling_ifeval.png" width="100%"><br>
+    </td>
+  </tr>
+</table>
+<figcaption align="left">
+<sub><b>Figure 6: Scaling of Optimal Learning Rate (&eta;) with Global Batch Size (B).</b> <i>For these experiments, the global batch size is parameterized by the total number of tokens processed per optimization step. Blue points represent empirical findings from Qwen3 training runs. The red dashed line shows a linear regression fit in log-log space, indicating a power-law relationship.</i></sub>
+</figcaption>
+</figure>
+
+To quantify this scaling relationship, we conducted a systematic hyperparameter sweep to identify the optimal learning rate across a broad spectrum of batch sizes. Analysis of these optimal pairings (illustrated in [Figure 6](#fig-lr-scale-gbs)) reveals a consistent logarithmic trend. Utilizing a least-squares fit, we derived a practical scaling law for production environments:
+
+$$\log(\eta) = a \log(B) + b$$
+
+For our specific configuration, we found <span>$$a = 0.578$$</span> and <span>$$b = -15.684$$</span>. By exponentiating both sides, we can express the learning rate as a power function of the batch size:
+
+$$\eta = B^a \cdot e^b$$
+
+**The Scaling Factor:**
+This relationship allows us to predict how the learning rate should change when the global batch size is scaled by a factor of <span>$$C$$</span>. If we define <span>$$\eta'$$</span> as the new learning rate for a scaled batch size <span>$$(C \cdot B)$$</span>, the ratio of the new learning rate to the original is:
+
+$$\frac{\eta'}{\eta} = \frac{(C \cdot B)^a \cdot e^b}{B^a \cdot e^b} = \left(\frac{C \cdot B}{B}\right)^a = C^a$$
+
+**Practical Takeaway:**
+This derivation provides a reliable heuristic for scaling your training runs on VTC. Essentially, when you scale your global batch size by <span>$$C$$</span>, you should scale your learning rate by <span>$$C^{0.578}$$</span>.
+
+**Example:** If you double your global batch size (<span>$$C = 2$$</span>), your learning rate should increase by a factor of <span>$$2^{0.578} \approx 1.49$$</span>. This "1.5x rule" ensures that your model remains on the optimal convergence path even as you significantly increase compute throughput.
+
+### Scaling learning rates based on model parameters
+
+While scaling batch size helps with throughput, another fundamental factor influencing your learning rate is the scale of the model itself. As we transition from compact edge models to large-scale dense architectures, the optimal learning rate (<span>$$\eta$$</span>) shifts predictably.
+
+<figure align="center" id="fig-lr-vs-val-loss">
+<table align="center" width="80%">
+  <tr>
+    <td align="center" width="100%">
+      <img src="images/lr_scaling_model_size.png" width="100%"><br>
+    </td>
+  </tr>
+</table>
+<figcaption align="left">
+<sub><b>Figure 7: Empirical Learning Rate Sweep across Model Scales.</b> <i>Validation loss is plotted against learning rate for five model sizes ranging from 0.6B to 14B parameters (embedding parameters are removed from the model size calculation). Stars indicate the observed minima for each configuration.</i></sub>
+</figcaption>
+</figure>
+
+To map this shift for the Qwen3 dense model family, we evaluated five distinct model scales: 0.6B, 1.7B, 4B, 8B, and 14B parameters. For each architecture, we performed a log-scale grid search to pinpoint the optimal learning rate (illustrated in [Figure 7](#fig-lr-vs-val-loss)). By applying a least-squares fit to these empirical data points, we established a power-law relationship between model parameters (<span>$$N$$</span>) and the learning rate ([Figure 8](#fig-lr-scale-model-size)):
+
+$$\log(\eta) = -0.646 \log(N) + 4.133$$
+
+**Scaling by Model Size:**
+Following a similar derivation to our batch size analysis, this formula allows us to predict the necessary adjustment when increasing model capacity. If the number of model parameters (<span>$$N$$</span>) increases by a factor of <span>$$C$$</span>, the optimal learning rate should be scaled by <span>$$C^{-0.646}$$</span>.
+
+**Practical Takeaway:**
+This inverse relationship means that as your model grows larger, your learning rate must become more conservative to maintain stability.
+
+**Example:** If you decide to double your model size (<span>$$C = 2$$</span>), the optimal learning rate should be multiplied by <span>$$2^{-0.646} \approx 0.639$$</span>.
+
+<figure align="center" id="fig-lr-scale-model-size">
+<table align="center" width="80%">
+  <tr>
+    <td align="center" width="100%">
+      <img src="images/model_scaling_ifeval.png" width="100%"><br>
+    </td>
+  </tr>
+</table>
+<figcaption align="left">
+<sub><b>Figure 8: Relationship between Optimal Learning Rate (&eta;) and Model Parameters (N).</b> <i>Empirical data points (blue) represent the best-performing learning rates across a parameter range of approximately **0.6B to 14B**. The red dashed line depicts the power-law scaling trend. The negative slope indicates that as parameter count increases, the learning rate must be scaled down according to a fixed ratio to maintain training efficiency.</i></sub>
+</figcaption>
+</figure>
+
+### Scaling learning rates based on token budget
+
+This section examines the relationship between optimal learning rates and the total training token budget. In contrast to the power-law relationships observed in pre-training literature (e.g., the [Chinchilla scaling laws](https://arxiv.org/abs/2203.15556)), our empirical findings indicate that the optimal learning rate remains stable as the token budget increases (see [Figure 9](#fig-lr-vs-token-budget)). This divergence suggests that the hyperparameter dynamics of Supervised Fine-Tuning (SFT) differ from those of initial pre-training phases.
+
+<figure align="center" id="fig-lr-vs-token-budget">
+<table align="center" width="80%">
+  <tr>
+    <td align="center" width="100%">
+      <img src="images/lr_scaling_token_budget.png" width="100%"><br>
+    </td>
+  </tr>
+</table>
+<figcaption align="left">
+<sub><b>Figure 9: Scaling of Optimal Learning Rate (&eta;) with Total Token Budget (T).</b> <i>Empirical data showing the relationship between the learning rate and the number of training tokens.</i></sub>
+</figcaption>
+</figure>
+
+**Practical Takeaway:**
+Within an SFT framework, the optimal learning rate is largely invariant to changes in the training token budget, allowing for consistent hyperparameter application across varying dataset scales.
+
+## Key Takeaways
+
+Thanks for reading. We hope this distillation framework and these scaling insights help you achieve frontier-level performance for your own reasoning models on Vertex AI Training Cluster.
+
+### Distillation Methodology & Teacher Selection
+
+* **The "Capacity Matching" Nuance:** Bigger is not always better. For well-defined reasoning tasks (Mathematics and IFEval), **capacity-matched** (same-sized) teachers often outperform larger models. However, for novel or highly complex domains like ARC-AGI, massive teacher models remain superior.
+* **Rollout Volume Matters:** Increasing the number of teacher rollouts per prompt captures the teacher’s inherent uncertainty and multiple valid reasoning paths. Performance generally increases with rollout count, though gains may saturate depending on the domain (e.g., IFEval saturates at 8 rollouts).
+* **Rejection Sampling is Domain-Specific:** While rejection sampling (training only on correct answers) provides a significant uplift for **Instruction Following**, it does not yield gains in Mathematics, ARC-AGI, or <span>$$\tau^2$$</span>-bench. In complex reasoning domains, "near-miss" cases appear essential for building robustness.
+
+### Hyperparameter Scaling
+
+The blog establishes three critical "rules of thumb" for scaling Supervised Fine-Tuning (SFT) hyperparameters:
+
+* **Batch Size Scaling:** When scaling the Global Batch Size (<span>$$B$$</span>) by a factor of <span>$$C$$</span>, the learning rate (<span>$$\eta$$</span>) should be scaled by <span>$$C^{0.578}$$</span>. For example, doubling the batch size suggests a **1.5x increase** in the learning rate.
+* **Model Size Scaling:** As model parameters (<span>$$N$$</span>) increase, the learning rate must become more conservative. The optimal learning rate scales by <span>$$C^{-0.646}$$</span> when the model size is increased by factor <span>$$C$$</span>.
+* **Token Budget Stability:** Unlike initial pre-training, the optimal learning rate for SFT is **largely invariant** to the total training token budget, allowing for consistent application across different dataset scales.
+
+## Acknowledgements
+
+We would like to express our sincere gratitude to the NVIDIA NeMo RL team–specifically Terry Kong– as well as the NVIDIA NeMo Gym team–specifically Brian Yu and Chris Wing– for their invaluable support throughout this project.
+
+We would also like to express our gratitude to our VTC teammates: Mohammadreza Mohseni, Mayank Sharan, Weiran Zhao, Jiuqiang Tang, Bo Wu, Lav Rai, and Minwoo Park for their infrastructure support, feedback, and insightful discussions throughout the project. We also thank Ting Yu, Shengyang Dai, Peng Xu, and Saurabh Tiwary for their leadership and support.

notebooks/community/CODEOWNERS

@@ -27,7 +27,9 @@
 /vertex_endpoints/nvidia-triton/nvidia-triton-custom-container-prediction.ipynb @RajeshThallam
 /vertex_endpoints/optimized_tensorflow_runtime @vlasenkoalexey
 /notebooks/community/alphagenome/cloudai_alphagenome_vai_quickstart.ipynb @dpanigra
+/notebooks/community/alphagenome/cloudai_alphagenome_finetune.ipynb @dpanigra
 /notebooks/community/weathernext/weathernext_2_early_access_program.ipynb @dpanigra
+/notebooks/community/weathernext/weathernext_2_ic_early_access_program.ipynb @dpanigra
 /notebooks/community/ml_ops/stage2/get_started_with_visionapi_and_automl.ipynb @mansari
 /notebooks/community/neo4j/graph_paysim.ipynb @benofben @laeg
 /notebooks/community/ml_ops/stage1/get_started_with_visionapi_and_vertex_datasets.ipynb @mansari

notebooks/community/alphagenome/README.md

@@ -1,7 +1,8 @@
 ![AlphaGenome header image](https://raw.githubusercontent.com/google-deepmind/alphagenome/refs/heads/main/docs/source/_static/header.png)
 
 # AlphaGenome
-[**Overview**](#overview) | [**Use Cases**](#use-cases) | [**Documentation**](#documentation) | [**Pricing**](#pricing) | [**Quick start**](#quick-start)
+[**Overview**](#overview) | [**Use Cases**](#use-cases) | [**Documentation**](#documentation) | [**Pricing**](#pricing) | [**Quick start inference**](#quick-start-inference) | 
+[**Quick start finetune**](#quick-start-finetune)
 
 ## Overview
 **Disclaimer:** *Experimental*.
@@ -89,5 +90,8 @@ To utilize these models via this service:
 *   **Pricing information** will be shared directly with users upon approval
     and placement on the allowlist.
 
-## Quick start
+## Quick start inference
 The quickest way to get started with the AlphaGenome in Google Cloud Platform is to run [our example notebook](cloudai_alphagenome_vai_quickstart.ipynb) in [Google Colab](https://colab.research.google.com/).
+
+## Quick start finetune
+The quickest way to get started with the AlphaGenome fineutning in Google Cloud Platform is to run [our finetuning notebook](cloudai_alphagenome_finetune.ipynb) in [Google Cloud Platform Enterprise Colab](https://docs.cloud.google.com/colab/docs/introduction).

notebooks/community/model_garden/docker_source_codes/model_oss/notebook_util/common_util.py

@@ -233,7 +233,7 @@ def download_image(url: str) -> str:
     base64 encoded image.
   """
   response = requests.get(url)
-  return Image.open(io.BytesIO(response.content))  # pytype: disable=bad-return-type  # pillow-102-upgrade
+  return Image.open(io.BytesIO(response.content))
 
 
 def resize_image(image: Any, new_width: int = 1000) -> Any:

notebooks/community/model_garden/docker_source_codes/model_oss/notebook_util/dataset_validation_util.py

@@ -387,6 +387,57 @@ def load_tokenizer(
   return tokenizer
 
 
+def _get_indices_for_valid_length(
+    dataset: Any,
+    input_column: str,
+    max_sequence_length: int,
+    tokenizer: transformers.PreTrainedTokenizer,
+    context_name: str = "the dataset",
+) -> tuple[list[int], int, int]:
+  """Gets indices of examples shorter than or equal to max_seq_length.
+
+  Args:
+    dataset: The dataset to check.
+    input_column: The input column in the dataset.
+    max_sequence_length: The maximum sequence length.
+    tokenizer: The tokenizer.
+    context_name: A name for the dataset used in log messages.
+
+  Returns:
+    A tuple of (indices_to_keep, original_length, dropped_samples).
+  """
+  if not dataset:
+    return [], 0, 0
+
+  original_length = len(dataset)
+  indices_to_keep = [
+      i
+      for i, entry in enumerate(dataset)
+      if len(tokenizer(entry[input_column])["input_ids"]) <= max_sequence_length
+  ]
+  dropped_samples = original_length - len(indices_to_keep)
+
+  if dropped_samples > 0:
+    examples_removed_percent = (dropped_samples * 100) / original_length
+    logging.info(
+        "(%.2f%%) of examples token length is <= max-seq-length(%d); (%.2f%%) >"
+        " max-seq-length in %s. %d example(s) were longer than max-seq-length.",
+        100 - examples_removed_percent,
+        max_sequence_length,
+        examples_removed_percent,
+        context_name,
+        dropped_samples,
+    )
+  else:
+    logging.info(
+        "No samples were dropped from %s because all samples are"
+        " shorter than max_sequence_length=%d.",
+        context_name,
+        max_sequence_length,
+    )
+  return indices_to_keep, original_length, dropped_samples
+
+
 def get_filtered_dataset(
     dataset: Any,
     input_column: str,
@@ -411,33 +462,25 @@ def get_filtered_dataset(
     ValueError: If more than `example_removed_threshold` of the dataset is
       filtered out.
   """
-  actual_dataset_length = len(dataset)
-  filtered_dataset = dataset.filter(
-      lambda x: len(tokenizer(x[input_column])["input_ids"]) <= max_seq_length
-  )
-  filtered_dataset_length = len(filtered_dataset)
-  if actual_dataset_length != filtered_dataset_length:
-    examples_removed_percent = (
-        (actual_dataset_length - filtered_dataset_length)
-        * 100
-        / actual_dataset_length
-    )
-    logging.info(
-        "(%.2f%%) of examples token length is <= max-seq-length(%d); (%.2f%%) >"
-        " max-seq-length. Filtering out %d example(s) which are longer than"
-        " max-seq-length.",
-        100 - examples_removed_percent,
-        max_seq_length,
-        examples_removed_percent,
-        actual_dataset_length - filtered_dataset_length,
-    )
-    if examples_removed_percent > example_removed_threshold:
-      raise ValueError(
-          "More than %.2f%% of the dataset is filtered out. This may be due to"
-          " small value of max-seq-length(%d) or incorrect template. Please"
-          " increase the max-seq-length or check the template."
-          % (examples_removed_percent, max_seq_length)
+  indices_to_keep, original_length, dropped_samples = (
+      _get_indices_for_valid_length(
+          dataset, input_column, max_seq_length, tokenizer, "the dataset"
       )
+  )
+
+  if (
+      original_length > 0
+      and dropped_samples / original_length * 100 > example_removed_threshold
+  ):
+    examples_removed_percent = (dropped_samples * 100) / original_length
+    raise ValueError(
+        f"More than {examples_removed_percent:.2f}% of the dataset is filtered"
+        " out. This may be due to small value of"
+        f" max-seq-length({max_seq_length}) or incorrect template. Please"
+        " increase the max-seq-length or check the template."
+    )
+
+  filtered_dataset = dataset.select(indices_to_keep)
   print(f"Some formatted examples from the dataset are: {filtered_dataset[:5]}")
   return filtered_dataset
 
@@ -502,6 +545,43 @@ def load_dataset_with_template(
   return raw, templated
 
 
+def drop_long_sequences(
+    dataset: Any,
+    dataset_with_template: Any,
+    input_column: str,
+    max_sequence_length: int,
+    tokenizer: transformers.PreTrainedTokenizer,
+    is_train: bool,
+) -> tuple[Any, Any, int]:
+  """Drops examples longer than max_seq_length from the dataset.
+
+  Args:
+    dataset: The dataset to filter.
+    dataset_with_template: The dataset with template to filter.
+    input_column: The input column in the dataset to be used.
+    max_sequence_length: The maximum sequence length.
+    tokenizer: The tokenizer.
+    is_train: Whether the dataset is for training.
+
+  Returns:
+    A tuple of (filtered_dataset, filtered_dataset_with_template,
+    dropped_samples).
+  """
+
+  context_name = f"the {'train' if is_train else 'eval'} dataset"
+  indices_to_keep, _, dropped_samples = _get_indices_for_valid_length(
+      dataset_with_template,
+      input_column,
+      max_sequence_length,
+      tokenizer,
+      context_name,
+  )
+
+  filtered_dataset = dataset.select(indices_to_keep)
+  filtered_dataset_with_template = dataset_with_template.select(indices_to_keep)
+  return filtered_dataset, filtered_dataset_with_template, dropped_samples
+
+
 def validate_dataset_with_template(
     dataset_name: str,
     split: str,

notebooks/community/model_garden/docker_source_codes/model_oss/timesfm/dockerfile/serve.Dockerfile

@@ -1,58 +0,0 @@
-FROM nvidia/cuda:12.3.2-devel-ubuntu22.04
-
-# Install basic libs
-RUN apt-get update && apt-get upgrade -y && apt-get install -y --no-install-recommends \
-        cmake \
-        curl \
-        wget \
-        sudo \
-        gnupg \
-        libsm6 \
-        libxext6 \
-        libxrender-dev \
-        lsb-release \
-        ca-certificates \
-        build-essential \
-        git \
-        software-properties-common \
-        cuda-toolkit \
-        libcudnn8 \
-        apt-transport-https
-
-RUN apt install -y --no-install-recommends python3.10 \
-        python3.10-venv \
-        python3.10-dev \
-        python3-pip
-
-Run apt-get autoremove -y
-
-RUN pip install --upgrade pip
-RUN pip install --upgrade --ignore-installed \
-        "jax[cuda12]==0.4.26" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html \
-        numpy==1.26.4 \
-        paxml==1.4.0 \
-        praxis==1.4.0 \
-        jaxlib==0.4.26 \
-        pandas==2.1.4 \
-        einshape==1.0.0 \
-        utilsforecast==0.1.10 \
-        huggingface_hub[cli]==0.23.0 \
-        google-cloud-aiplatform[prediction]==1.51.0 \
-        fastapi==0.109.1 \
-        flask==3.0.3 \
-        smart_open[gcs]==7.0.4 \
-        protobuf==3.19.6 \
-        scikit-learn==1.0.2 \
-        timesfm==1.0.1
-
-# Download license.
-RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
-
-# Move scaffold.
-COPY model_oss/timesfm/main.py /app/main.py
-COPY model_oss/timesfm/predictor.py /app/predictor.py
-
-WORKDIR ..
-
-# Spin off inference server.
-CMD ["python3", "/app/main.py"]

notebooks/community/model_garden/docker_source_codes/model_oss/timesfm/main.py

@@ -1,71 +0,0 @@
-"""Predict server for TimesFM."""
-
-import json
-import os
-import flask
-import predictor
-from predictor import PredictionError
-
-# Create the flask app.
-app = flask.Flask(__name__)
-_OK_STATUS = 200
-_INTERNAL_ERROR_STATUS = 500
-_BAD_REQUEST_STATUS = 400
-_HOST = '0.0.0.0'
-
-# Define the predictor and load the checkpoints.
-predictor = predictor.TimesFMPredictor()
-predictor.load(os.environ['AIP_STORAGE_URI'])
-
-
-@app.route(os.environ['AIP_HEALTH_ROUTE'], methods=['GET'])
-def health() -> flask.Response:
-  return flask.Response(status=_OK_STATUS)
-
-
-@app.route(os.environ['AIP_PREDICT_ROUTE'], methods=['GET', 'POST'])
-def predict() -> flask.Response:
-  """Calls TimesFM for prediction.
-
-  Returns:
-    A `flask.Response` containing the prediction result in JSON.
-  """
-  try:
-    body = flask.request.get_json(silent=True, force=True)
-    preprocessed_inputs = predictor.preprocess(body)
-    outputs = predictor.predict(preprocessed_inputs)
-    conf_level = preprocessed_inputs.get('conf_level')
-
-    if conf_level is not None:
-      postprocessed_outputs = predictor.postprocess_with_conf_level(
-          outputs, preprocessed_inputs['conf_level']
-      )
-    else:
-      postprocessed_outputs = predictor.postprocess(outputs)
-    return flask.Response(
-        json.dumps(postprocessed_outputs),
-        status=_OK_STATUS,
-        mimetype='application/json',
-    )
-  except PredictionError as e:
-    return flask.Response(
-        json.dumps({'error': str(e)}),
-        status=e.status_code,
-        mimetype='application/json',
-    )
-  except ValueError as e:
-    return flask.Response(
-        json.dumps({'error': str(e)}),
-        status=_BAD_REQUEST_STATUS,
-        mimetype='application/json',
-    )
-  except Exception as e:  # pylint: disable=broad-exception-caught
-    return flask.Response(
-        json.dumps({'error': str(e)}),
-        status=_INTERNAL_ERROR_STATUS,
-        mimetype='application/json',
-    )
-
-
-if __name__ == '__main__':
-  app.run(host=_HOST, port=os.environ['AIP_HTTP_PORT'])

notebooks/community/model_garden/docker_source_codes/model_oss/timesfm/predictor.py

@@ -1,609 +0,0 @@
-"""Adapts a pretrained TimesFM to the CPR framework.
-
-Documentation for the model is here:
-https://github.com/google-research/timesfm
-
-Model checkpoints can be found here:
-https://www.huggingface.co/google/timesfm-1.0-200m
-"""
-
-from collections.abc import Sequence
-import datetime
-import os
-from typing import Any
-import fastapi
-from google.cloud.aiplatform.utils import prediction_utils
-from jax._src import config
-import numpy as np
-import scipy.stats as st
-import timesfm
-
-
-HTTPException = fastapi.HTTPException
-_BACKEND = os.getenv("TIMESFM_BACKEND", default="cpu")
-
-config.update(
-    "jax_platforms", {"cpu": "cpu", "gpu": "cuda", "tpu": ""}[_BACKEND]
-)
-
-TsArray = None | float | int | str | list["TsArray"]
-
-_BAD_REQUEST_STATUS = 400
-_EXPECTED_FORMAT = """
-
-[NOTICE] TimesFM inference server expects input format:
-{
-    "instances": [
-        {
-            "input": [0.0, 0.1, 0.2, ...],
-            "freq": 0,                                       # optional, 0/1/2
-            "horizon": 12,                                   # optional
-            "timestamp": ["2024-01-01", "2024-01-02", ...],  # optional
-            "timestamp_format": "%Y-%m-%d",                  # optional
-            "dynamic_numerical_covariates": {
-                "dncov1": [1.0, 2.0, 1.5, ...],
-                "dncov2": [3.0, 1.1, 2.4, ...],
-            },                                               # optional
-            "dynamic_categorical_covariates": {
-                "dccov1": ["a", "b", "a", ...],
-                "dccov2": [0, 1, 0, ...],
-            },                                               # optional
-            "static_numerical_covariates": {
-                "sncov1": 1.0,
-                "sncov2": 2.0,
-            },                                               # optional
-            "static_categorical_covariates": {
-                "sccov1": "a",
-                "sccov2": "b",
-            },                                               # optional
-            "xreg_kwargs": {...},                            # optional
-        },
-        {"input": [113.2, 15.0, 65.4, ...], ...},
-        {"input": [  0.0, 10.0, 20.0, ...], ...},
-        ...
-    ]
-}
-"""
-
-
-class PredictionError(Exception):
-  """Custom exception for prediction errors."""
-
-  def __init__(self, message: str, status_code: int = _BAD_REQUEST_STATUS):
-    super().__init__(message)
-    self.status_code = status_code
-    self.message = message
-
-
-def _raise_bad_request(message: str):
-  message = message + "\n" + _EXPECTED_FORMAT
-  raise PredictionError(
-      message=message,
-      status_code=_BAD_REQUEST_STATUS,
-  )
-
-
-def _datetime_to_freq(dt1: datetime.datetime, dt2: datetime.datetime) -> int:
-  delta = dt2 - dt1
-  if delta.days <= 1:
-    return 0
-  elif delta.days <= 31:
-    return 1
-  else:
-    return 2
-
-
-def _add_cov_to_dict(
-    index: int,
-    cov_input: dict[str, TsArray],
-    cov_dict: dict[str, list[TsArray]],
-):
-  """Adds covariates to the dictionary of covariates.
-
-  Args:
-    index: Index of the instance.
-    cov_input: Dictionary of covariates for the current instance.
-    cov_dict: Dictionary of covariates for all instances.
-  """
-  if index == 0:
-    cov_dict.update({k: [v] for k, v in cov_input.items()})
-  else:
-    if set(cov_input.keys()) != set(cov_dict.keys()):
-      _raise_bad_request(
-          f"Instance {index}:"
-          " All instances must have the same set of covariates if any."
-      )
-    for k, v in cov_input.items():
-      cov_dict[k].append(v)
-
-
-def _linear_interpolate_missing_timepoints(
-    timestamp: list[datetime.datetime],
-    value: list[float],
-) -> tuple[list[datetime.datetime], list[TsArray]]:
-  """Linearly interpolates missing timepoints in a timeseries."""
-
-  def _gcd_timelapse(t1, t2):
-    if (w := t2 % t1) == datetime.timedelta(0):
-      return t1
-    if t1 > t2:
-      return _gcd_timelapse(t2, t1)
-    return _gcd_timelapse(w, t1)
-
-  if len(timestamp) < 3:
-    return timestamp, value, False
-
-  no_missing = True
-  delta = timestamp[1] - timestamp[0]
-  if delta <= datetime.timedelta(0):
-    _raise_bad_request(
-        f"Timestamps must be in ascending order. Got {timestamp}"
-    )
-  for i in range(2, len(timestamp)):
-    delta_next = timestamp[i] - timestamp[i - 1]
-    if delta_next <= datetime.timedelta(0):
-      _raise_bad_request(
-          f"Timestamps must be in ascending order. Got {timestamp}"
-      )
-    delta_new = _gcd_timelapse(delta, delta_next)
-    if delta_new != delta:
-      no_missing = False
-      delta = delta_new
-
-  if no_missing:
-    return timestamp, value, False
-
-  new_timestamp = []
-  new_value = []
-  for i in range(len(timestamp) - 1):
-    new_timestamp.append(timestamp[i])
-    new_value.append(value[i])
-    if (num_deltas := int((timestamp[i + 1] - timestamp[i]) / delta + 0.5)) > 1:
-      value_delta = (value[i + 1] - value[i]) / num_deltas
-      for j in range(1, num_deltas):
-        new_timestamp.append(timestamp[i] + j * delta)
-        new_value.append(value[i] + j * value_delta)
-  new_timestamp.append(timestamp[-1])
-  new_value.append(value[-1])
-
-  return new_timestamp, new_value, True
-
-
-class TimesFMPredictor:
-  """Predictor class for time-series foundation model TimesFM."""
-
-  TIMESFM_MODEL_NAME = os.getenv(
-      "TIMESFM_MODEL_NAME", default="timesfm-1.0-200m"
-  )
-  CONTEXT_LEN = 512
-  INPUT_PATCH_LEN = 32
-  OUTPUT_PATCH_LEN = 128
-  NUM_LAYERS = 20
-  MODEL_DIMS = 1280
-  BACKEND = os.getenv("TIMESFM_BACKEND", default="cpu")
-  MAX_HORIZON = int(os.getenv("TIMESFM_HORIZON", default="128"))
-
-  def load(self, artifacts_uri: str = ""):
-    """Initializes the model and preprocessing transforms.
-
-    Args:
-      artifacts_uri: Directory where state dict is stored. Can be a GCS URI or
-        local path.
-    """
-    if not (os.path.isdir(artifacts_uri) or artifacts_uri.startswith("gs://")):
-      raise ValueError(
-          f"Provided artifact_uri is not a directory: {artifacts_uri}"
-      )
-
-    print(f"Downloading checkpoints from {artifacts_uri}")
-    prediction_utils.download_model_artifacts(artifacts_uri)
-    artifact_path = os.getcwd()
-
-    print(f"Loading checkpoints from {artifact_path}")
-    self._model = timesfm.TimesFm(
-        context_len=self.CONTEXT_LEN,
-        horizon_len=(
-            ((self.MAX_HORIZON - 1) // self.OUTPUT_PATCH_LEN + 1)
-            * self.OUTPUT_PATCH_LEN
-        ),
-        input_patch_len=self.INPUT_PATCH_LEN,
-        output_patch_len=self.OUTPUT_PATCH_LEN,
-        num_layers=self.NUM_LAYERS,
-        model_dims=self.MODEL_DIMS,
-        backend=self.BACKEND,
-    )
-    self._model.load_from_checkpoint(artifact_path)
-    print(f"Loaded TimesFM model from {artifact_path}")
-
-  def preprocess(
-      self, request_dict: dict[str, Sequence[dict[str, TsArray]]]
-  ) -> dict[str, TsArray]:
-    """Performs preprocessing.
-
-    By default, the server expects a request body consisting of a valid JSON
-    object. This will be parsed by the handler before it's evaluated by the
-    preprocess method.
-
-    Args:
-      request_dict: Parsed request body. We expect that the input consists of a
-        list of time-series forecast contexts. Each context should be in a
-        format convertible to JTensor by `jnp.array`.
-
-    Returns:
-      Time-series forecast contexts are passed as is from the input as a list.
-    """
-
-    if "instances" not in request_dict:
-      _raise_bad_request('Request must contain "instances" as a top-level key.')
-
-    input_instances = request_dict["instances"]
-    if not input_instances or not isinstance(input_instances, list):
-      _raise_bad_request(
-          f"Received `instances` not a list. Got {type(input_instances)}"
-      )
-
-    inputs, freqs, timestamps, timestamp_formats = [], [], [], []
-    horizon_lens = []
-    conf_level = None
-    static_numerical_covariates, static_categorical_covariates = {}, {}
-    dynamic_numerical_covariates, dynamic_categorical_covariates = {}, {}
-    xreg_kwargs = {}
-
-    exists_missing = False
-    for index, each_input in enumerate(input_instances):
-
-      # 1. Add input time-series context.
-      if (
-          (not isinstance(each_input, dict))
-          or ("input" not in each_input)
-          or (len(each_input["input"]) < 2)
-      ):
-        _raise_bad_request(
-            f"Instance {index}:"
-            " Invalid datatype. Each input example must have `input` key"
-            " mapped to a list of time-series forecast context with length > 1."
-        )
-      new_input = each_input["input"]
-
-      # 2. Process timestamps.
-      if "timestamp" not in each_input:
-        timestamps.append(None)
-      else:
-        if len(each_input["timestamp"]) != len(each_input["input"]):
-          _raise_bad_request(
-              f"Instance {index}:"
-              " Invalid datatype. `timestamp` if given must have same length as"
-              "`input`."
-          )
-        new_timestamp = [
-            datetime.datetime.fromisoformat(s) for s in each_input["timestamp"]
-        ]
-        # Linearly interpolate missing timepoints and values.
-        new_timestamp, new_input, new_exists_missing = (
-            _linear_interpolate_missing_timepoints(new_timestamp, new_input)
-        )
-        exists_missing = exists_missing or new_exists_missing
-        timestamps.append(new_timestamp)
-      if "timestamp_format" in each_input:
-        timestamp_formats.append(each_input["timestamp_format"])
-      else:
-        timestamp_formats.append(None)
-
-      inputs.append(new_input)
-
-      # 3. Process frequency.
-      if "freq" in each_input:
-        freqs.append(each_input["freq"])
-      elif timestamps[index]:
-        freqs.append(
-            _datetime_to_freq(timestamps[index][0], timestamps[index][1])
-        )
-      else:
-        freqs.append(0)
-
-      # 4. Process covariate data.
-      for cov_category, cov_dict in [
-          ("static_numerical_covariates", static_numerical_covariates),
-          ("static_categorical_covariates", static_categorical_covariates),
-          ("dynamic_numerical_covariates", dynamic_numerical_covariates),
-          ("dynamic_categorical_covariates", dynamic_categorical_covariates),
-      ]:
-        if cov_category in each_input:
-          _add_cov_to_dict(index, each_input[cov_category], cov_dict)
-
-      # 5. Process xreg config. Power user option. If nothing set we apply
-      # TimesFM default.
-      if "xreg_kwargs" in each_input:
-        if not xreg_kwargs:
-          xreg_kwargs = each_input["xreg_kwargs"]
-        elif xreg_kwargs != each_input["xreg_kwargs"]:
-          _raise_bad_request(
-              f"Instance {index}:"
-              " All instances must have the same xreg_kwargs if any."
-          )
-
-      # 6. Process horizon length.
-      if "horizon" in each_input:
-        if (w := each_input["horizon"]) > self.MAX_HORIZON:
-          _raise_bad_request(
-              f"Instance {index}: `horizon` must be <= maximum horizon"
-              f" {self.MAX_HORIZON}. Got {w}. To increase the maximum horizon,"
-              " recreate the endpoint with a higher `TIMESFM_HORIZON` env"
-              " value."
-          )
-        horizon_lens.append(w)
-      else:
-        horizon_lens.append(self.MAX_HORIZON)
-
-      # 7. Process conf level.
-      all_conf_levels = [
-          each_input.get("conf_level", None) for each_input in input_instances
-      ]
-
-      defined_conf_levels = [cl for cl in all_conf_levels if cl is not None]
-      undefined_conf_levels = [cl for cl in all_conf_levels if cl is None]
-
-      if defined_conf_levels and undefined_conf_levels:
-        _raise_bad_request(
-            "Either all or none of the instances must define `conf_level`."
-        )
-
-      if defined_conf_levels:
-        unique_conf_levels = set(defined_conf_levels)
-        if len(unique_conf_levels) > 1:
-          _raise_bad_request("All instances must have the same `conf_level`.")
-        conf_level = unique_conf_levels.pop()
-        if not 0 <= conf_level <= 1:
-          _raise_bad_request(
-              f"`conf_level` must be between 0 and 1. Got {conf_level}."
-          )
-      else:
-        conf_level = None
-
-    return {
-        "inputs": inputs,
-        "freqs": freqs,
-        "timestamps": timestamps,
-        "timestamp_formats": timestamp_formats,
-        "exists_missing": exists_missing,
-        "static_numerical_covariates": static_numerical_covariates,
-        "static_categorical_covariates": static_categorical_covariates,
-        "dynamic_numerical_covariates": dynamic_numerical_covariates,
-        "dynamic_categorical_covariates": dynamic_categorical_covariates,
-        "xreg_kwargs": xreg_kwargs,
-        "horizon_lens": horizon_lens,
-        "conf_level": conf_level,
-    }
-
-  def predict(self, instances: dict[str, Any]) -> Any:
-    """Performs prediction.
-
-    Args:
-      instances: A dictionary with two keys - `inputs` and `freq` where `inputs`
-        is list of time series forecast contexts. Each context time series
-        should be in a format convertible to JTensor by `jnp.array`. `freq` is
-        frequencies of each forecast context with values as 0 (high), 1 (medium)
-        and 2 (low). If not provided, all contexts are assumed to be high
-        frequency.
-
-    Returns:
-        A tuple of List:
-        - the mean forecast of size (# inputs, # forecast horizon),
-        - the full forecast (mean + quantiles) of size
-            (# inputs,  # forecast horizon, 1 + # quantiles).
-    """
-    (
-        inputs,
-        freqs,
-        timestamps,
-        timestamp_formats,
-        exists_missing,
-        static_numerical_covariates,
-        static_categorical_covariates,
-        dynamic_numerical_covariates,
-        dynamic_categorical_covariates,
-        xreg_kwargs,
-        horizon_lens,
-    ) = (
-        instances["inputs"],
-        instances["freqs"],
-        instances["timestamps"],
-        instances["timestamp_formats"],
-        instances["exists_missing"],
-        instances["static_numerical_covariates"],
-        instances["static_categorical_covariates"],
-        instances["dynamic_numerical_covariates"],
-        instances["dynamic_categorical_covariates"],
-        instances["xreg_kwargs"],
-        instances["horizon_lens"],
-    )
-
-    if (
-        static_numerical_covariates
-        or static_categorical_covariates
-        or dynamic_numerical_covariates
-        or dynamic_categorical_covariates
-    ):
-      if (
-          dynamic_categorical_covariates or dynamic_numerical_covariates
-      ) and exists_missing:
-        _raise_bad_request(
-            "Dynamic covariates are not supported when input has missing"
-            " timestamps."
-        )
-      print("Detected covariates. Callng model.forecast_with_covariates.")
-      try:
-        point_forecast, _ = self._model.forecast_with_covariates(
-            inputs=inputs,
-            dynamic_numerical_covariates=dynamic_numerical_covariates,
-            dynamic_categorical_covariates=dynamic_categorical_covariates,
-            static_numerical_covariates=static_numerical_covariates,
-            static_categorical_covariates=static_categorical_covariates,
-            freq=freqs,
-            **xreg_kwargs,
-        )
-        # point_forecast is a list of np.ndarrays.
-        point_forecast = [p.tolist() for p in point_forecast]
-        quantile_forecast = None
-      except ValueError as e:
-        _raise_bad_request(f"model.forecast_with_covariates failed from {e}.")
-        return
-    else:
-      print("Calling model.forecast.")
-      point_forecast, quantile_forecast = self._model.forecast(
-          inputs=inputs, freq=freqs
-      )
-      # point_forecast and quantile_forecast are JTensors (np.ndarrays).
-      point_forecast = point_forecast.tolist()
-      quantile_forecast = quantile_forecast.tolist()
-
-    return (
-        point_forecast,
-        quantile_forecast,
-        timestamps,
-        timestamp_formats,
-        horizon_lens,
-    )
-
-  def postprocess(
-      self, forecasts: tuple[TsArray, TsArray, TsArray, TsArray]
-  ) -> dict[str, list[dict[str, TsArray]]]:
-    """Translates the model output.
-
-    Args:
-       forecasts: A tuple of List - the mean forecast of size (# inputs, #
-         forecast horizon), - the full forecast (mean + quantiles) of size (#
-         inputs,  # forecast horizon, 1 + # quantiles).
-
-    Returns:
-      Dictionary containing the list of point forecasts and quantile forecasts
-      for each of the input time-series context.
-    """
-    (
-        point_forecasts,
-        quantile_forecasts,
-        timestamps,
-        timestamp_formats,
-        horizon_lens,
-    ) = forecasts
-    predictions = []
-
-    quantile_names = ["mean"] + [
-        f"p{int(quantile * 100)}" for quantile in self._model.model_p.quantiles
-    ]
-    for i, point_forecast in enumerate(point_forecasts):
-      response = {"point_forecast": point_forecast[: horizon_lens[i]]}
-
-      if quantile_forecasts:
-        for j, quantile_name in enumerate(quantile_names):
-          response[quantile_name] = [x[j] for x in quantile_forecasts[i]][
-              : horizon_lens[i]
-          ]
-
-      if timestamps[i]:
-        last_timestamp = timestamps[i][-1]
-        timestamp_delta = timestamps[i][-1] - timestamps[i][-2]
-        response["timestamp"] = []
-        for _ in range(len(point_forecast)):
-          last_timestamp = last_timestamp + timestamp_delta
-          response["timestamp"].append(
-              datetime.datetime.strftime(last_timestamp, timestamp_formats[i])
-              if timestamp_formats[i]
-              else last_timestamp.isoformat()
-          )
-        response["timestamp"] = response["timestamp"][: horizon_lens[i]]
-
-      predictions.append(response)
-    return {"predictions": predictions}
-
-  def postprocess_with_conf_level(
-      self,
-      forecasts: tuple[TsArray, TsArray, TsArray, TsArray, TsArray],
-      conf_level: float | None,
-  ) -> dict[str, list[dict[str, TsArray]]]:
-    """Translates the model output."""
-    lower_quantile = (1 - conf_level) / 2
-    higher_quantile = (1 + conf_level) / 2
-    _, quantile_forecast, _, _, horizon_lens = forecasts
-    response = self.postprocess(forecasts)
-
-    if quantile_forecast is None:
-      return response
-
-    # Note: The raw quantile forecast from TimesFM has the mean as the 0-th
-    # element. We strip it before passing to extend_quantiles.
-    quantile_forecast_np_array = np.array(quantile_forecast)
-    extended_forecasts = extend_quantiles(
-        quantile_forecast_np_array[..., 1:],
-        lower_quantile,
-        higher_quantile,
-        model_quantiles=self._model.model_p.quantiles,
-    )
-    lower_bounds = extended_forecasts["lower_bound"]
-    upper_bounds = extended_forecasts["upper_bound"]
-
-    for i, prediction in enumerate(response["predictions"]):
-      horizon = horizon_lens[i]
-      prediction["lower_bound"] = lower_bounds[i][:horizon].tolist()
-      prediction["upper_bound"] = upper_bounds[i][:horizon].tolist()
-
-    return response
-
-
-def extend_quantiles(
-    quantile_forecast: np.ndarray,
-    lower_quantile: float,
-    higher_quantile: float,
-    model_quantiles: list[float],
-) -> dict[str, np.ndarray]:
-  """Extends the quantile forecast to the lower and upper bounds.
-
-  Args:
-    quantile_forecast: The quantile forecast from TimesFM.
-    lower_quantile: The lower quantile to extend to.
-    higher_quantile: The higher quantile to extend to.
-    model_quantiles: The quantiles used by the model.
-
-  Returns:
-    A dictionary containing the lower and upper bounds.
-  """
-  if quantile_forecast.shape[2] != len(model_quantiles):
-    raise ValueError(
-        "Number of model quantiles should match the last dimension of the"
-        " quantile forecast. If you are using the raw TimesFM quantile forecast"
-        "output, you likely need to strip the 0-index which is the mean."
-    )
-
-  idx_median = model_quantiles.index(0.5)
-  idx_low_q = np.argmin(model_quantiles)
-  low_q = model_quantiles[idx_low_q]
-
-  if not (low_q < 0.5):
-    raise ValueError(
-        f"The lowest quantile {low_q=} provided in the forecast must be less"
-        " than 0.5."
-    )
-  idx_high_q = np.argmax(model_quantiles)
-  high_q = model_quantiles[idx_high_q]
-
-  if not (high_q > 0.5):
-    raise ValueError(
-        f"The highest quantile {high_q=} provided in the forecast must be"
-        " greater than 0.5."
-    )
-
-  positive_sigma = np.maximum(
-      0, quantile_forecast[..., idx_high_q] - quantile_forecast[..., idx_median]
-  ) / st.norm.ppf(high_q)
-  negative_sigma = np.minimum(
-      0, quantile_forecast[..., idx_low_q] - quantile_forecast[..., idx_median]
-  ) / st.norm.ppf(low_q)
-
-  lower_bound = quantile_forecast[
-      ..., idx_median
-  ] + negative_sigma * st.norm.ppf(lower_quantile)
-  upper_bound = quantile_forecast[
-      ..., idx_median
-  ] + positive_sigma * st.norm.ppf(higher_quantile)
-
-  return {"lower_bound": lower_bound, "upper_bound": upper_bound}

notebooks/community/model_garden/docker_source_codes/notebook_util/common_util.py

@@ -233,7 +233,7 @@ def download_image(url: str) -> str:
     base64 encoded image.
   """
   response = requests.get(url)
-  return Image.open(io.BytesIO(response.content))  # pytype: disable=bad-return-type  # pillow-102-upgrade
+  return Image.open(io.BytesIO(response.content))
 
 
 def resize_image(image: Any, new_width: int = 1000) -> Any:

notebooks/community/model_garden/docker_source_codes/notebook_util/dataset_validation_util.py

@@ -387,6 +387,57 @@ def load_tokenizer(
   return tokenizer
 
 
+def _get_indices_for_valid_length(
+    dataset: Any,
+    input_column: str,
+    max_sequence_length: int,
+    tokenizer: transformers.PreTrainedTokenizer,
+    context_name: str = "the dataset",
+) -> tuple[list[int], int, int]:
+  """Gets indices of examples shorter than or equal to max_seq_length.
+
+  Args:
+    dataset: The dataset to check.
+    input_column: The input column in the dataset.
+    max_sequence_length: The maximum sequence length.
+    tokenizer: The tokenizer.
+    context_name: A name for the dataset used in log messages.
+
+  Returns:
+    A tuple of (indices_to_keep, original_length, dropped_samples).
+  """
+  if not dataset:
+    return [], 0, 0
+
+  original_length = len(dataset)
+  indices_to_keep = [
+      i
+      for i, entry in enumerate(dataset)
+      if len(tokenizer(entry[input_column])["input_ids"]) <= max_sequence_length
+  ]
+  dropped_samples = original_length - len(indices_to_keep)
+
+  if dropped_samples > 0:
+    examples_removed_percent = (dropped_samples * 100) / original_length
+    logging.info(
+        "(%.2f%%) of examples token length is <= max-seq-length(%d); (%.2f%%) >"
+        " max-seq-length in %s. %d example(s) were longer than max-seq-length.",
+        100 - examples_removed_percent,
+        max_sequence_length,
+        examples_removed_percent,
+        context_name,
+        dropped_samples,
+    )
+  else:
+    logging.info(
+        "No samples were dropped from %s because all samples are"
+        " shorter than max_sequence_length=%d.",
+        context_name,
+        max_sequence_length,
+    )
+  return indices_to_keep, original_length, dropped_samples
+
+
 def get_filtered_dataset(
     dataset: Any,
     input_column: str,
@@ -411,33 +462,25 @@ def get_filtered_dataset(
     ValueError: If more than `example_removed_threshold` of the dataset is
       filtered out.
   """
-  actual_dataset_length = len(dataset)
-  filtered_dataset = dataset.filter(
-      lambda x: len(tokenizer(x[input_column])["input_ids"]) <= max_seq_length
-  )
-  filtered_dataset_length = len(filtered_dataset)
-  if actual_dataset_length != filtered_dataset_length:
-    examples_removed_percent = (
-        (actual_dataset_length - filtered_dataset_length)
-        * 100
-        / actual_dataset_length
-    )
-    logging.info(
-        "(%.2f%%) of examples token length is <= max-seq-length(%d); (%.2f%%) >"
-        " max-seq-length. Filtering out %d example(s) which are longer than"
-        " max-seq-length.",
-        100 - examples_removed_percent,
-        max_seq_length,
-        examples_removed_percent,
-        actual_dataset_length - filtered_dataset_length,
-    )
-    if examples_removed_percent > example_removed_threshold:
-      raise ValueError(
-          "More than %.2f%% of the dataset is filtered out. This may be due to"
-          " small value of max-seq-length(%d) or incorrect template. Please"
-          " increase the max-seq-length or check the template."
-          % (examples_removed_percent, max_seq_length)
+  indices_to_keep, original_length, dropped_samples = (
+      _get_indices_for_valid_length(
+          dataset, input_column, max_seq_length, tokenizer, "the dataset"
       )
+  )
+
+  if (
+      original_length > 0
+      and dropped_samples / original_length * 100 > example_removed_threshold
+  ):
+    examples_removed_percent = (dropped_samples * 100) / original_length
+    raise ValueError(
+        f"More than {examples_removed_percent:.2f}% of the dataset is filtered"
+        " out. This may be due to small value of"
+        f" max-seq-length({max_seq_length}) or incorrect template. Please"
+        " increase the max-seq-length or check the template."
+    )
+
+  filtered_dataset = dataset.select(indices_to_keep)
   print(f"Some formatted examples from the dataset are: {filtered_dataset[:5]}")
   return filtered_dataset
 
@@ -502,6 +545,43 @@ def load_dataset_with_template(
   return raw, templated
 
 
+def drop_long_sequences(
+    dataset: Any,
+    dataset_with_template: Any,
+    input_column: str,
+    max_sequence_length: int,
+    tokenizer: transformers.PreTrainedTokenizer,
+    is_train: bool,
+) -> tuple[Any, Any, int]:
+  """Drops examples longer than max_seq_length from the dataset.
+
+  Args:
+    dataset: The dataset to filter.
+    dataset_with_template: The dataset with template to filter.
+    input_column: The input column in the dataset to be used.
+    max_sequence_length: The maximum sequence length.
+    tokenizer: The tokenizer.
+    is_train: Whether the dataset is for training.
+
+  Returns:
+    A tuple of (filtered_dataset, filtered_dataset_with_template,
+    dropped_samples).
+  """
+
+  context_name = f"the {'train' if is_train else 'eval'} dataset"
+  indices_to_keep, _, dropped_samples = _get_indices_for_valid_length(
+      dataset_with_template,
+      input_column,
+      max_sequence_length,
+      tokenizer,
+      context_name,
+  )
+
+  filtered_dataset = dataset.select(indices_to_keep)
+  filtered_dataset_with_template = dataset_with_template.select(indices_to_keep)
+  return filtered_dataset, filtered_dataset_with_template, dropped_samples
+
+
 def validate_dataset_with_template(
     dataset_name: str,
     split: str,

notebooks/community/model_garden/gke_model_ui_deployment_notebook.ipynb

@@ -10,7 +10,7 @@
       },
       "outputs": [],
       "source": [
-        "# Copyright 2025 Google LLC\n",
+        "# Copyright 2026 Google LLC\n",
         "#\n",
         "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
         "# you may not use this file except in compliance with the License.\n",

notebooks/community/model_garden/gke_model_ui_deployment_notebook_auto.ipynb

@@ -10,7 +10,7 @@
       },
       "outputs": [],
       "source": [
-        "# Copyright 2025 Google LLC\n",
+        "# Copyright 2026 Google LLC\n",
         "#\n",
         "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
         "# you may not use this file except in compliance with the License.\n",

notebooks/community/model_garden/model_garden_advanced_features.ipynb

@@ -6,11 +6,11 @@
       "id": "DZ1j6RRg-Td6",
       "metadata": {
         "cellView": "form",
-        "id": "f705f4be70e9"
+        "id": "483138c1a042"
       },
       "outputs": [],
       "source": [
-        "# Copyright 2025 Google LLC\n",
+        "# Copyright 2026 Google LLC\n",
         "#\n",
         "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
         "# you may not use this file except in compliance with the License.\n",

notebooks/community/model_garden/model_garden_axolotl_finetuning.ipynb

@@ -9,7 +9,7 @@
       },
       "outputs": [],
       "source": [
-        "# Copyright 2025 Google LLC\n",
+        "# Copyright 2026 Google LLC\n",
         "#\n",
         "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
         "# you may not use this file except in compliance with the License.\n",

notebooks/community/model_garden/model_garden_axolotl_gemma3_finetuning.ipynb

@@ -9,7 +9,7 @@
       },
       "outputs": [],
       "source": [
-        "# Copyright 2025 Google LLC\n",
+        "# Copyright 2026 Google LLC\n",
         "#\n",
         "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
         "# you may not use this file except in compliance with the License.\n",

notebooks/community/model_garden/model_garden_axolotl_gpt_oss_finetuning.ipynb

@@ -9,7 +9,7 @@
       },
       "outputs": [],
       "source": [
-        "# Copyright 2025 Google LLC\n",
+        "# Copyright 2026 Google LLC\n",
         "#\n",
         "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
         "# you may not use this file except in compliance with the License.\n",

notebooks/community/model_garden/model_garden_axolotl_qwen3_finetuning.ipynb

@@ -9,7 +9,7 @@
       },
       "outputs": [],
       "source": [
-        "# Copyright 2025 Google LLC\n",
+        "# Copyright 2026 Google LLC\n",
         "#\n",
         "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
         "# you may not use this file except in compliance with the License.\n",

notebooks/community/model_garden/model_garden_codegemma_deployment_on_vertex.ipynb

@@ -9,7 +9,7 @@
       },
       "outputs": [],
       "source": [
-        "# Copyright 2025 Google LLC\n",
+        "# Copyright 2026 Google LLC\n",
         "#\n",
         "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
         "# you may not use this file except in compliance with the License.\n",

notebooks/community/model_garden/model_garden_deployment_tutorial.ipynb

@@ -9,7 +9,7 @@
       },
       "outputs": [],
       "source": [
-        "# Copyright 2025 Google LLC\n",
+        "# Copyright 2026 Google LLC\n",
         "#\n",
         "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
         "# you may not use this file except in compliance with the License.\n",

notebooks/community/model_garden/model_garden_e5.ipynb

@@ -9,7 +9,7 @@
       },
       "outputs": [],
       "source": [
-        "# Copyright 2025 Google LLC\n",
+        "# Copyright 2026 Google LLC\n",
         "#\n",
         "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
         "# you may not use this file except in compliance with the License.\n",