google-cloud-pipeline-components 2.6.0__py3-none-any.whl → 2.7.0__py3-none-any.whl

google_cloud_pipeline_components/_implementation/llm/arbiter_preprocess.py

@@ -0,0 +1,137 @@
+# Copyright 2023 The Kubeflow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""KFP Container component for preprocessing predictions for the Arbiter."""
+
+import os
+from typing import Dict, List
+
+from google_cloud_pipeline_components import _placeholders
+from google_cloud_pipeline_components import utils as gcpc_utils
+from google_cloud_pipeline_components._implementation.llm import utils
+from kfp import dsl
+
+
+def _resolve_image() -> str:
+  """Determines the image URI to create a container from."""
+  return (
+      os.environ.get('AUTOSXS_IMAGE_OVERRIDE')
+      or utils.get_default_image_uri('autosxs'))
+
+
+# pylint: disable=unused-argument,dangerous-default-value
+@dsl.container_component
+def arbiter_preprocess(
+    evaluation_dataset: str,
+    id_columns: List[str],
+    response_column_a: str,
+    response_column_b: str,
+    task: str,
+    is_bp_output_a: bool,
+    is_bp_output_b: bool,
+    autorater_prompt_parameters: Dict[str, Dict[str, str]],
+    preprocessed_evaluation_dataset: dsl.Output[dsl.Dataset],  # pylint: disable=unused-argument # pytype: disable=unsupported-operands
+    preprocessed_evaluation_dataset_uri: dsl.OutputPath(str),  # pylint: disable=unused-argument # pytype: disable=invalid-annotation
+    gcp_resources: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    prediction_uris_a: str = '',
+    prediction_uris_b: str = '',
+    model_a_prompt_parameters: Dict[str, Dict[str, str]] = {},
+    model_b_prompt_parameters: Dict[str, Dict[str, str]] = {},
+    human_preference_column: str = '',
+) -> dsl.ContainerSpec:  # pylint: disable=g-doc-args
+  """Preprocesses predictions tables for the AutoSxS Arbiter.
+
+  Args:
+    evaluation_dataset: GCS or BigQuery URIs representing a dataset of prompts
+      and responses.
+    id_columns: The columns which distinguish unique evaluation examples.
+    response_column_a: The column containing responses for model a.
+    response_column_b: The column containing responses for model a.
+    task: Task to evaluate.
+    output_path: Path to write the path where preprocessed predictions are
+      stored.
+    is_bp_output_a: If True, the prediction URIs will be parsed as if they came
+      from Vertex Batch Prediction, where response_column_a represents a field
+      in the model output containing the response. If False, the expected format
+      will be a table containing all model_prompt_parameters and the
+      response_column.
+    is_bp_output_b: If True, the prediction URIs will be parsed as if they came
+      from Vertex Batch Prediction, where response_column_b represents a field
+      in the model output containing the response. If False, the expected format
+      will be a table containing all model_prompt_parameters and the
+      response_column.
+    prediction_uris: A list of GCS or BigQuery URIs representing a dataset of
+      prompts and responses for model a.
+    prediction_uris: A list of GCS or BigQuery URIs representing a dataset of
+      prompts and responses for model b.
+    model_a_prompt_parameters: Map of model A prompt template parameters to
+      columns or templates.
+    model_b_prompt_parameters: Map of model B prompt template parameters to
+      columns or templates.
+    autorater_prompt_parameters: Map of autorater prompt template parameters to
+      columns or templates.
+    human_preference_column: The column containing ground truths. The default
+      value is an empty string if not be provided by users.
+
+  Returns:
+    preprocessed_evaluation_dataset: Dataset of the table containing the inputs
+    expected by the Arbiter.
+    preprocessed_evaluation_dataset_uri: URI of the table containing the inputs
+    expected by the Arbiter.
+    gcp_resources: Tracker for GCP resources created by this component.
+  """
+  return gcpc_utils.build_serverless_customjob_container_spec(
+      project=_placeholders.PROJECT_ID_PLACEHOLDER,
+      location=_placeholders.LOCATION_PLACEHOLDER,
+      custom_job_payload=utils.build_payload(
+          display_name='arbiter_preprocess',
+          machine_type='n1-standard-4',
+          image_uri=_resolve_image(),
+          args=[
+              '--',  # Used to mark the start of component flags.
+              'arbiter_preprocess',
+              f'--evaluation_dataset={evaluation_dataset}',
+              f'--prediction_uris_a={prediction_uris_a}',
+              f'--prediction_uris_b={prediction_uris_b}',
+              (
+                  '--id_columns='
+                  "{{$.inputs.parameters['id_columns'].json_escape[0]}}"
+              ),
+              (
+                  '--autorater_prompt_parameters='
+                  "{{$.inputs.parameters['autorater_prompt_parameters']"
+                  '.json_escape[0]}}'
+              ),
+              (
+                  '--model_a_prompt_parameters='
+                  "{{$.inputs.parameters['model_a_prompt_parameters']"
+                  '.json_escape[0]}}'
+              ),
+              (
+                  '--model_b_prompt_parameters='
+                  "{{$.inputs.parameters['model_b_prompt_parameters']"
+                  '.json_escape[0]}}'
+              ),
+              f'--response_column_a={response_column_a}',
+              f'--response_column_b={response_column_b}',
+              f'--human_preference_column={human_preference_column}',
+              f'--task={task}',
+              f'--is_batch_prediction_output_a={is_bp_output_a}',
+              f'--is_batch_prediction_output_b={is_bp_output_b}',
+              f'--output_dir={dsl.PIPELINE_ROOT_PLACEHOLDER}',
+              f'--preprocessed_evaluation_dataset_uri={preprocessed_evaluation_dataset_uri}',
+              '--executor_input={{$.json_escape[1]}}',
+          ],
+      ),
+      gcp_resources=gcp_resources,
+  )

google_cloud_pipeline_components/_implementation/llm/autosxs_arbiter.py

@@ -0,0 +1,105 @@
+# Copyright 2023 The Kubeflow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""KFP Container component that performs AutoSxS."""
+
+import os
+from typing import Any, Dict, List
+
+from google_cloud_pipeline_components import _placeholders
+from google_cloud_pipeline_components import utils as gcpc_utils
+from google_cloud_pipeline_components._implementation.llm import utils
+from kfp import dsl
+
+
+def _resolve_image() -> str:
+  """Determines the image URI to create a container from."""
+  return (
+      os.environ.get('AUTOSXS_IMAGE_OVERRIDE')
+      or utils.get_default_image_uri('autosxs'))
+
+
+def _get_prediction_endpoint_overrides() -> str:
+  """Used for integration tests to override the prediction endpoint."""
+  return os.environ.get('PREDICTION_ENDPOINT_OVERRIDES', '')
+
+
+@dsl.container_component
+def autosxs_arbiter(
+    inference_output_uri: str,
+    id_columns: List[str],
+    task: str,
+    judgments: dsl.Output[dsl.Dataset],  # pylint: disable=unused-argument # pytype: disable=unsupported-operands
+    judgments_uri: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    gcp_resources: dsl.OutputPath(str),
+    metadata: dsl.OutputPath(str),
+    human_preference_column: str = '',
+    judgments_format: str = 'jsonl',
+    bigquery_destination_prefix: str = '',
+    experimental_args: Dict[str, Any] = {},
+) -> dsl.ContainerSpec:  # pylint: disable=g-doc-args
+  """Evaluate two models using an autorater.
+
+  Args:
+    inference_output_uri: Directory of model A's inference output.
+    id_columns: The columns which distinguish unique evaluation examples.
+    human_preference_column: Human preference column included in our inference
+      output.
+    task: Evaluation task in the form {task}@{version}. task can be one of
+      "summarization", "question_answer". Version is an integer with 3 digits or
+      "latest". Ex: summarization@001 or question_answer@latest.
+    judgments_format: The format to write judgments to. Can be either 'json' or
+      'bigquery'.
+    bigquery_destination_prefix: BigQuery table to write judgments to if the
+      specified format is 'bigquery'.
+    experimental_args: Experimentally released arguments. Subject to change.
+
+  Returns:
+    judgments: Individual judgments used to calculate the win rates.
+    judgments_uri: URI of the Judgments Artifact.
+    gcp_resources: Tracker for GCP resources created by this component.
+    metadata: Computed runtime metrics metadata from this component.
+  """
+  return gcpc_utils.build_serverless_customjob_container_spec(
+      project=_placeholders.PROJECT_ID_PLACEHOLDER,
+      # Hardcode location to us-central1 for text-bison availability.
+      location='us-central1',
+      custom_job_payload=utils.build_payload(
+          display_name='autosxs_arbiter',
+          machine_type='n1-standard-4',
+          image_uri=_resolve_image(),
+          args=[
+              '--',  # Used to mark the start of component flags.
+              'arbiter',
+              f'--inference_output_uri={inference_output_uri}',
+              f'--human_preference_column={human_preference_column}',
+              f'--task={task}',
+              f'--prediction_endpoint_overrides={_get_prediction_endpoint_overrides()}',
+              f'--output_dir={dsl.PIPELINE_ROOT_PLACEHOLDER}',
+              f'--judgments_uri={judgments_uri}',
+              f'--judgments_format={judgments_format}',
+              f'--bigquery_destination_prefix={bigquery_destination_prefix}',
+              (
+                  '--id_columns='
+                  "{{$.inputs.parameters['id_columns'].json_escape[0]}}"
+              ),
+              (
+                  '--experimental_args='
+                  "{{$.inputs.parameters['experimental_args'].json_escape[0]}}"
+              ),
+              '--executor_input={{$.json_escape[1]}}',
+              f'--metadata_path={metadata}',
+          ],
+      ),
+      gcp_resources=gcp_resources,
+  )

google_cloud_pipeline_components/_implementation/llm/autosxs_metrics_computer.py

@@ -0,0 +1,66 @@
+# Copyright 2023 The Kubeflow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""KFP Container component for computing AutoSXS metrics."""
+
+import os
+
+from google_cloud_pipeline_components import _placeholders
+from google_cloud_pipeline_components import utils as gcpc_utils
+from google_cloud_pipeline_components._implementation.llm import utils
+from kfp import dsl
+
+
+def _resolve_image() -> str:
+  """Determines the image URI to create a container from."""
+  return os.environ.get(
+      'AUTOSXS_IMAGE_OVERRIDE'
+  ) or utils.get_default_image_uri('autosxs')
+
+
+@dsl.container_component
+def autosxs_metrics_computer(
+    judgments_dir: str,
+    has_human_preference: bool,
+    autosxs_metrics: dsl.Output[dsl.Metrics],  # pylint: disable=unused-argument # pytype: disable=unsupported-operands
+    gcp_resources: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+) -> dsl.ContainerSpec:  # pylint: disable=g-doc-args
+  """Compute AutoSXS metrics using judgments outputs from Arbiter.
+
+  Args:
+    judgments_dir: Path where store the Judgments.
+    has_human_preference: Boolean value. True if users provided human preference
+      data, otherwise false.
+
+  Returns:
+    autosxs_metrics: Autosxs win rate metrics and human alignment metrics.
+    gcp_resources: Tracker for GCP resources created by this component.
+  """
+  return gcpc_utils.build_serverless_customjob_container_spec(
+      project=_placeholders.PROJECT_ID_PLACEHOLDER,
+      # Hardcode location to us-central1 for text-bison availability.
+      location='us-central1',
+      custom_job_payload=utils.build_payload(
+          display_name='autosxs_metrics_computer',
+          machine_type='n1-standard-4',
+          image_uri=_resolve_image(),
+          args=[
+              '--',  # Used to mark the start of component flags.
+              'autosxs_metrics',
+              f'--judgments_dir={judgments_dir}',
+              f'--has_human_preference={has_human_preference}',
+              '--executor_input={{$.json_escape[1]}}',
+          ],
+      ),
+      gcp_resources=gcp_resources,
+  )

google_cloud_pipeline_components/_implementation/llm/deployment_graph.py

@@ -75,22 +75,16 @@ def pipeline(
           'large_model_reference'
       ]
   ).set_display_name('Resolve Upload Model')
-  upload_task = (
-      upload_llm_model.upload_llm_model(
-          project=_placeholders.PROJECT_ID_PLACEHOLDER,
-          location=upload_location,
-          regional_endpoint=regional_endpoint.output,
-          artifact_uri=adapter_artifact.output,
-          model_display_name=display_name.output,
-          model_reference_name='text-bison@001',
-          upload_model=upload_model.output,
-      )
-      .set_env_variable(
-          name='VERTEX_AI_PIPELINES_RUN_LABELS',
-          value=json.dumps({'tune-type': 'rlhf'}),
-      )
-      .set_display_name('Upload Model')
-  )
+  upload_task = upload_llm_model.upload_llm_model(
+      project=_placeholders.PROJECT_ID_PLACEHOLDER,
+      location=upload_location,
+      regional_endpoint=regional_endpoint.output,
+      artifact_uri=adapter_artifact.output,
+      model_display_name=display_name.output,
+      model_reference_name='text-bison@001',
+      upload_model=upload_model.output,
+      tune_type='rlhf',
+  ).set_display_name('Upload Model')
   deploy_model = function_based.resolve_deploy_model(
       deploy_model=deploy_model,
       large_model_reference=reference_model_metadata.outputs[

google_cloud_pipeline_components/_implementation/llm/env.py

@@ -16,7 +16,7 @@ import os
 
 
 def get_private_image_tag() -> str:
-  return os.getenv('PRIVATE_IMAGE_TAG', '20231031_0507_RC00')
+  return os.getenv('PRIVATE_IMAGE_TAG', '20231213_0507_RC00')
 
 
 def get_use_test_machine_spec() -> bool:

google_cloud_pipeline_components/_implementation/llm/function_based.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 """Python function-based components used in KFP pipelies."""
 import functools
-from typing import List, NamedTuple, Optional
+from typing import Any, Dict, List, NamedTuple, Optional
 
 from google_cloud_pipeline_components import _image
 from google_cloud_pipeline_components._implementation.llm import env
@@ -302,8 +302,8 @@ def resolve_reference_model_metadata(
       'llama-2-13b': reference_model_metadata(
           large_model_reference='LLAMA_2_13B',
           reference_model_path='gs://vertex-rlhf-restricted/pretrained_models/llama/t5x_llama_2_13b/',
-          reward_model_reference='LLAMA_2_13B',
-          reward_model_path='gs://vertex-rlhf-restricted/pretrained_models/llama/t5x_llama_2_13b/',
+          reward_model_reference='LLAMA_2_7B',
+          reward_model_path='gs://vertex-rlhf-restricted/pretrained_models/llama/t5x_llama_2_7b/',
           is_supported=True,
       ),
       'llama-2-7b-chat': reference_model_metadata(
@@ -316,8 +316,8 @@ def resolve_reference_model_metadata(
       'llama-2-13b-chat': reference_model_metadata(
           large_model_reference='LLAMA_2_13B_CHAT',
           reference_model_path='gs://vertex-rlhf-restricted/pretrained_models/llama/t5x_llama_2_13b_chat/',
-          reward_model_reference='LLAMA_2_13B',
-          reward_model_path='gs://vertex-rlhf-restricted/pretrained_models/llama/t5x_llama_2_13b/',
+          reward_model_reference='LLAMA_2_7B',
+          reward_model_path='gs://vertex-rlhf-restricted/pretrained_models/llama/t5x_llama_2_7b/',
           is_supported=True,
       ),
   }
@@ -495,3 +495,80 @@ def resolve_instruction(
   """
   instruction = instruction or ''
   return instruction if 'chat' not in large_model_reference.lower() else ''
+
+
+@dsl.component(base_image=_image.GCPC_IMAGE_TAG, install_kfp_package=False)
+def resolve_num_microbatches(large_model_reference: str) -> int:
+  """Resolves the number of microbatches to use during training.
+
+  Args:
+    large_model_reference: Base model tuned by the pipeline.
+
+  Returns:
+    Number of microbatches to break the total batch size into during training.
+  """
+  if 'llama' in large_model_reference.lower():
+    return 2
+  return 0
+
+
+@dsl.component(base_image=_image.GCPC_IMAGE_TAG, install_kfp_package=False)
+def read_file(path: str) -> str:
+  """Reads the contents of the given file."""
+  # pylint: disable=g-import-not-at-top,import-outside-toplevel,redefined-outer-name,reimported
+  import re
+  # pylint: enable=g-import-not-at-top,import-outside-toplevel,redefined-outer-name,reimported
+
+  path = re.sub('^gs://', '/gcs/', path)
+  with open(path, 'r') as f:
+    return f.read()
+
+
+@dsl.component(base_image=_image.GCPC_IMAGE_TAG, install_kfp_package=False)
+def get_usage_metric(metadata: Dict[str, Any], key: str) -> bool:  # pytype: disable=unsupported-operands
+  """Extracts a single usage metric from metadata."""
+  return metadata[key]
+
+
+@dsl.component(base_image=_image.GCPC_IMAGE_TAG, install_kfp_package=False)
+def dump_dict(value: Dict[Any, Any]) -> str:
+  """Dumps the given dict to a JSON string."""
+  # pylint: disable=g-import-not-at-top,import-outside-toplevel,redefined-outer-name,reimported
+  import json
+  # pylint: enable=g-import-not-at-top,import-outside-toplevel,redefined-outer-name,reimported
+
+  return json.dumps(value).replace('"', '\\"')
+
+
+@dsl.component(base_image=_image.GCPC_IMAGE_TAG, install_kfp_package=False)
+def dump_list(value: List[Any]) -> str:
+  """Dumps the given dict to a JSON string."""
+  # pylint: disable=g-import-not-at-top,import-outside-toplevel,redefined-outer-name,reimported
+  import json
+  # pylint: enable=g-import-not-at-top,import-outside-toplevel,redefined-outer-name,reimported
+
+  return json.dumps(value).replace('"', '\\"')
+
+
+@dsl.component(base_image=_image.GCPC_IMAGE_TAG, install_kfp_package=False)
+def identity(
+    x: str,
+) -> str:
+  return x
+
+
+@dsl.component(base_image=_image.GCPC_IMAGE_TAG, install_kfp_package=False)
+def get_uri(artifact: dsl.Input[dsl.Artifact], is_dir: bool = False) -> str:  # pytype: disable=unsupported-operands
+  """Extracts the URI from an artifact."""
+  # pylint: disable=g-import-not-at-top,import-outside-toplevel,redefined-outer-name,reimported
+  import os
+  # pylint: enable=g-import-not-at-top,import-outside-toplevel,redefined-outer-name,reimported
+
+  if is_dir:
+    return os.path.join(artifact.uri, '*')
+  return artifact.uri
+
+
+@dsl.component(base_image=_image.GCPC_IMAGE_TAG, install_kfp_package=False)
+def get_empty_string() -> str:
+  return ''

google_cloud_pipeline_components/_implementation/llm/reinforcement_learning_graph.py

@@ -117,6 +117,11 @@ def pipeline(
       accelerator_type=machine_spec.outputs['accelerator_type'],
       accelerator_count=machine_spec.outputs['accelerator_count'],
   ).set_display_name('Resolve Reinforcer Image URI')
+  num_microbatches = function_based.resolve_num_microbatches(
+      large_model_reference=reference_model_metadata.outputs[
+          'large_model_reference'
+      ]
+  ).set_display_name('Resolve Number of Microbatches')
   rl_model = (
       reinforcer.Reinforcer(
           project=project,
@@ -145,6 +150,7 @@ def pipeline(
           learning_rate_multiplier=reinforcement_learning_rate_multiplier,
           kl_coeff=kl_coeff,
           lora_dim=lora_dim,
+          num_microbatches=num_microbatches.output,
       )
       .set_display_name('Reinforcer')
       .set_caching_options(False)

google_cloud_pipeline_components/_implementation/llm/reinforcer.py

@@ -43,6 +43,7 @@ def Reinforcer(  # pylint: disable=invalid-name
     learning_rate_multiplier: float = 1.0,
     kl_coeff: float = 0.1,
     lora_dim: int = 0,
+    num_microbatches: int = 0,
 ) -> kfp.dsl.ContainerSpec:  # pylint: disable=g-doc-args
   """Trains a model using reinforcement learning.
 
@@ -53,8 +54,8 @@ def Reinforcer(  # pylint: disable=invalid-name
     input_reward_model_path: Path to the reward model to use during
       reinforcement learning.
     input_dataset_path: Path to training dataset.
-    train_steps: Number of training steps. These are the number of steps
-      on top of any steps used to train the base model.
+    train_steps: Number of training steps. These are the number of steps on top
+      of any steps used to train the base model.
     targets_length: Maximum decoder steps. Outputs will be at most this length.
     accelerator_type: Type of TPU accelerator. Can be either TPU_V2 or TPU_V3.
     accelerator_count: Number of TPU accelerators.
@@ -75,6 +76,9 @@ def Reinforcer(  # pylint: disable=invalid-name
       then use full-tuning.
     learning_rate_multiplier: Constant multiplied by the base learning rate used
       to adjust the learning rate during reinforcement learning.
+    num_microbatches: Number of microbatches to break the total batch size into
+      during training. If <= 1, the model is trained on the full batch size
+      directly.
 
   Returns:
     output_model_path: Path to the trained model checkpoint.
@@ -110,6 +114,7 @@ def Reinforcer(  # pylint: disable=invalid-name
               f'--learning_rate_multiplier={learning_rate_multiplier}',
               f'--kl_coeff={kl_coeff}',
               f'--lora_dim={lora_dim}',
+              f'--num_microbatches={num_microbatches}',
           ],
       ),
       gcp_resources=gcp_resources,

google_cloud_pipeline_components/_implementation/llm/reward_model_graph.py

@@ -118,6 +118,11 @@ def pipeline(
       accelerator_type=machine_spec.outputs['accelerator_type'],
       accelerator_count=machine_spec.outputs['accelerator_count'],
   ).set_display_name('Resolve Reward Model Image URI')
+  num_microbatches = function_based.resolve_num_microbatches(
+      large_model_reference=reference_model_metadata.outputs[
+          'reward_model_reference'
+      ]
+  ).set_display_name('Resolve Number of Microbatches')
   reward_model = (
       reward_model_trainer.RewardModelTrainer(
           project=project,
@@ -141,6 +146,7 @@ def pipeline(
           batch_size=batch_size,
           learning_rate_multiplier=reward_model_learning_rate_multiplier,
           lora_dim=lora_dim,
+          num_microbatches=num_microbatches.output,
       )
       .set_display_name('Reward Model Trainer')
       .set_caching_options(False)

google_cloud_pipeline_components/_implementation/llm/reward_model_trainer.py

@@ -39,6 +39,7 @@ def RewardModelTrainer(  # pylint: disable=invalid-name
     batch_size: int = 64,
     learning_rate_multiplier: float = 1.0,
     lora_dim: int = 0,
+    num_microbatches: int = 0,
 ) -> kfp.dsl.ContainerSpec:  # pylint: disable=g-doc-args
   """Trains a reward model.
 
@@ -47,8 +48,8 @@ def RewardModelTrainer(  # pylint: disable=invalid-name
     location: Location used to run the job.
     input_model_path: Path to the base model to fine tune.
     input_dataset_path: Path to dataset to use to train a reward model.
-    train_steps: Number of training steps. These are the number of steps
-      on top of any steps used to train the base model.
+    train_steps: Number of training steps. These are the number of steps on top
+      of any steps used to train the base model.
     accelerator_type: Type of TPU accelerator. Can be either TPU_V2 or TPU_V3.
     accelerator_count: Number of TPU accelerators.
     large_model_reference: Predefined model used to create the ``input_model``.
@@ -64,6 +65,9 @@ def RewardModelTrainer(  # pylint: disable=invalid-name
       then use full-tuning.
     learning_rate_multiplier: Constant multiplied by the base learning rate used
       to adjust the learning rate when training a reward model.
+    num_microbatches: Number of microbatches to break the total batch size into
+      during training. If <= 1, the model is trained on the full batch size
+      directly.
 
   Returns:
     output_model: Trained reward model.
@@ -98,6 +102,7 @@ def RewardModelTrainer(  # pylint: disable=invalid-name
                   f'{kfp.dsl.PIPELINE_TASK_ID_PLACEHOLDER}'
               ),
               f'--lora_dim={lora_dim}',
+              f'--num_microbatches={num_microbatches}',
           ],
       ),
       gcp_resources=gcp_resources,

google_cloud_pipeline_components/_implementation/llm/supervised_fine_tuner.py

@@ -39,6 +39,7 @@ def SupervisedFineTuner(  # pylint: disable=invalid-name
     batch_size: int = 64,
     learning_rate_multiplier: float = 1.0,
     lora_dim: int = 0,
+    num_microbatches: int = 0,
 ) -> kfp.dsl.ContainerSpec:  # pylint: disable=g-doc-args
   """Performs supervised fine tuning.
 
@@ -65,6 +66,9 @@ def SupervisedFineTuner(  # pylint: disable=invalid-name
       then use full-tuning.
     learning_rate_multiplier: Constant multiplied by the base learning rate used
       to adjust the learning rate during supervised fine tuning.
+    num_microbatches: Number of microbatches to break the total batch size into
+      during training. If <= 1, the model is trained on the full batch size
+      directly.
 
   Returns:
     output_model_path: Fine-tuned model path.
@@ -99,6 +103,7 @@ def SupervisedFineTuner(  # pylint: disable=invalid-name
                   f'{kfp.dsl.PIPELINE_TASK_ID_PLACEHOLDER}'
               ),
               f'--lora_dim={lora_dim}',
+              f'--num_microbatches={num_microbatches}',
           ],
       ),
       gcp_resources=gcp_resources,