Fix gemma3 finetuning to improve model quality

PiperOrigin-RevId: 744724903
2026-09-26 14:42:04 +00:00 · 2025-04-07 08:08:56 -07:00
1 changed files with 2 additions and 4 deletions
@@ -426,9 +426,7 @@
        "        \"scheduling_strategy\": gca_custom_job_compat.Scheduling.Strategy.FLEX_START,\n",
        "    }\n",
        "\n",
-        "TRAIN_DOCKER_URI = (\n",
-        "    f\"{repo}/vertex-vision-model-garden-dockers/pytorch-peft-train:gemma3_20250312\"\n",
-        ")\n",
+        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/cloud-nas-260507/vertex-vision-model-garden-dockers/pytorch-peft-train:gemma3_07042025_5_30_ist\"\n",
        "\n",
        "# Worker pool spec.\n",
        "if training_accelerator_type == \"NVIDIA_A100_80GB\":\n",
@@ -464,7 +462,7 @@
        "# Attention implementation to use in the model.\n",
        "attn_implementation = \"eager\"\n",
        "# The optimizer for which to schedule the learning rate.\n",
-        "optimizer = \"paged_adamw_32bit\"\n",
+        "optimizer = \"adamw_torch_fused\"\n",
        "# Define the proportion of training to be dedicated to a linear warmup where learning rate gradually increases.\n",
        "warmup_ratio = \"0.01\"\n",
        "# The list or string of integrations to report the results and logs to.\n",