feat(components): add resolve_machine_spec and resolve_refined_image_…

…uri to rlhf_preprocessor component

PiperOrigin-RevId: 626080295

components/google-cloud/google_cloud_pipeline_components/_implementation/llm/generated/refined_image_versions.py

@@ -17,4 +17,4 @@
 DO NOT EDIT - This file is generated, manual changes will be overridden.
 """
 
-IMAGE_TAG = '20240414_0507'
+IMAGE_TAG = '20240417_0507_RC00'

components/google-cloud/google_cloud_pipeline_components/_implementation/llm/reinforcement_learning_graph.py

@@ -44,6 +44,11 @@ def pipeline(
     reward_model_reference: str,
     policy_model_reference: str,
     policy_model_path: str,
+    machine_type: str,
+    tuning_location: str,
+    accelerator_type: str,
+    accelerator_count: int,
+    rl_image_uri: str,
     prompt_sequence_length: int = 512,
     target_sequence_length: int = 64,
     lora_dim: int = 1,
@@ -54,7 +59,6 @@ def pipeline(
     kl_coeff: float = 0.1,
     instruction: Optional[str] = None,
     project: str = _placeholders.PROJECT_ID_PLACEHOLDER,
-    accelerator_type: str = 'GPU',
     location: str = _placeholders.LOCATION_PLACEHOLDER,
     tensorboard_resource_id: str = '',
     encryption_spec_key_name: str = '',
@@ -70,6 +74,11 @@ def pipeline(
     reward_model_reference: Name of the reward model. The name should be in capitalized snake case format.
     policy_model_reference: Name of the policy model. The name should be in capitalized snake case format.
     policy_model_path: The model checkpoint path to the reinforcer model.
+    machine_type: The type of the machine to provision for the custom job. Must be a valid GCE instance type and compatible with the accelerator type.
+    tuning_location: The GCP region to run the custom job.
+    accelerator_type: Specific accelerator type for the custom job.
+    accelerator_count: The number of accelerator.
+    rl_image_uri: Docker image URI to use for the reinforcement learning training job.
     prompt_sequence_length: Maximum tokenized sequence length for input text. Higher values increase memory overhead. This value should be at most 8192. Default value is 512.
     target_sequence_length: Maximum tokenized sequence length for target text. Higher values increase memory overhead. This value should be at most 1024. Default value is 64.
     lora_dim: The rank of the LoRA adapter. If >0, then use LoRA-tuning. If =0, then use full-tuning. Default is 1.
@@ -80,7 +89,6 @@ def pipeline(
     kl_coeff: Coefficient for KL penalty. This regularizes the policy model and penalizes if it diverges from its initial distribution. If set to 0, the reference language model is not loaded into memory. Default value is 0.1.
     instruction: This field lets the model know what task it needs to perform. Base models have been trained over a large set of varied instructions. You can give a simple and intuitive description of the task and the model will follow it, e.g. "Classify this movie review as positive or negative" or "Translate this sentence to Danish". Do not specify this if your dataset already prepends the instruction to the inputs field.
     project: Project used to run custom jobs. If not specified the project used to run the pipeline will be used.
-    accelerator_type: One of 'TPU' or 'GPU'. If 'TPU' is specified, tuning components run in europe-west4. Otherwise tuning components run in us-central1 on GPUs. Default is 'GPU'.
     location: Location used to run non-tuning components, i.e. components that do not require accelerators. If not specified the location used to run the pipeline will be used.
     tensorboard_resource_id: Optional tensorboard resource id in format `projects/{project_number}/locations/{location}/tensorboards/{tensorboard_id}`. If provided, tensorboard metrics will be uploaded to this location.
     encryption_spec_key_name: Customer-managed encryption key. If this is set, then all resources created by the CustomJob will be encrypted with the provided encryption key. Note that this is not supported for TPU at the moment.
@@ -91,10 +99,6 @@ def pipeline(
   """
   # fmt: on
   prompt_column = 'input_text'
-  machine_spec = function_based.resolve_machine_spec(
-      accelerator_type=accelerator_type,
-      use_test_spec=env.get_use_test_machine_spec(),
-  ).set_display_name('Resolve Machine Spec')
 
   processed_dataset = preprocess_chat_dataset.preprocess_chat_dataset(
       large_model_reference=large_model_reference,
@@ -118,16 +122,13 @@ def pipeline(
       .set_display_name('Import Prompt Dataset')
       .set_caching_options(False)
   )
-  rl_image_uri = function_based.resolve_private_refined_image_uri(
-      accelerator_type=machine_spec.outputs['accelerator_type'],
-  ).set_display_name('Resolve Reinforcer Image URI')
   num_microbatches = function_based.resolve_num_microbatches(
       large_model_reference=policy_model_reference,
   ).set_display_name('Resolve Number of Microbatches')
   rl_model = (
       reinforcer.reinforcer(
           project=project,
-          location=machine_spec.outputs['tuning_location'],
+          location=tuning_location,
           input_reference_model_path=policy_model_path,
           input_reward_model_path=input_reward_model_path,
           input_reward_adapter_path=input_reward_adapter_path,
@@ -136,12 +137,12 @@ def pipeline(
           ],
           input_preference_dataset_path=input_preference_dataset_path,
           train_steps=reinforcement_learning_train_steps,
-          accelerator_type=machine_spec.outputs['accelerator_type'],
-          accelerator_count=machine_spec.outputs['accelerator_count'],
+          accelerator_type=accelerator_type,
+          accelerator_count=accelerator_count,
           large_model_reference=policy_model_reference,
           reward_model_reference=reward_model_reference,
-          machine_type=machine_spec.outputs['machine_type'],
-          image_uri=rl_image_uri.output,
+          machine_type=machine_type,
+          image_uri=rl_image_uri,
           inputs_sequence_length=prompt_sequence_length,
           targets_sequence_length=target_sequence_length,
           batch_size=batch_size,

components/google-cloud/google_cloud_pipeline_components/_implementation/llm/reward_model_graph.py

@@ -40,6 +40,11 @@ def pipeline(
     large_model_reference: str,
     reward_model_reference: str,
     reward_model_path: str,
+    machine_type: str,
+    tuning_location: str,
+    accelerator_type: str,
+    accelerator_count: int,
+    reward_model_image_uri: str,
     prompt_sequence_length: int = 512,
     target_sequence_length: int = 64,
     batch_size: int = 64,
@@ -49,7 +54,6 @@ def pipeline(
     eval_dataset: Optional[str] = None,
     instruction: Optional[str] = None,
     project: str = _placeholders.PROJECT_ID_PLACEHOLDER,
-    accelerator_type: str = 'GPU',
     location: str = _placeholders.LOCATION_PLACEHOLDER,
     tensorboard_resource_id: str = '',
     encryption_spec_key_name: str = '',
@@ -62,6 +66,11 @@ def pipeline(
     large_model_reference: Name of the base model. Supported values are `text-bison@001`, `t5-small`, `t5-large`, `t5-xl` and `t5-xxl`. `text-bison@001` and `t5-small` are supported in `us-central1` and `europe-west4`. `t5-large`, `t5-xl` and `t5-xxl` are only supported in `europe-west4`.
     reward_model_reference: Name of the base model. The name should be in capitalized snake case format.
     reward_model_path: The model checkpoint path for the reward model.
+    machine_type: The type of the machine to provision for the custom job. Must be a valid GCE instance type and compatible with the accelerator type.
+    tuning_location: The GCP region to run the custom job.
+    accelerator_type: Specific accelerator type for the custom job.
+    accelerator_count: The number of accelerator.
+    reward_model_image_uri: Docker image URI to use for the reward model training job.
     prompt_sequence_length: Maximum tokenized sequence length for input text. Higher values increase memory overhead. This value should be at most 8192. Default value is 512.
     target_sequence_length:  Maximum tokenized sequence length for target text. Higher values increase memory overhead. This value should be at most 1024. Default value is 64.
     batch_size: Number of examples in each finetuning step. Default is 64.
@@ -70,7 +79,6 @@ def pipeline(
     reward_model_train_steps: Number of steps to use when training a reward model. Default value is 1000.
     instruction: This field lets the model know what task it needs to perform. Base models have been trained over a large set of varied instructions. You can give a simple and intuitive description of the task and the model will follow it, e.g. "Classify this movie review as positive or negative" or "Translate this sentence to Danish". Do not specify this if your dataset already prepends the instruction to the inputs field.
     project: Project used to run custom jobs. If not specified the project used to run the pipeline will be used.
-    accelerator_type: One of 'TPU' or 'GPU'. If 'TPU' is specified, tuning components run in europe-west4. Otherwise tuning components run in us-central1 on GPUs. Default is 'GPU'.
     location: Location used to run non-tuning components, i.e. components that do not require accelerators. If not specified the location used to run the pipeline will be used.
     tensorboard_resource_id: Optional tensorboard resource id in format `projects/{project_number}/locations/{location}/tensorboards/{tensorboard_id}`. If provided, tensorboard metrics will be uploaded to this location.
     encryption_spec_key_name: Customer-managed encryption key. If this is set, then all resources created by the CustomJob will be encrypted with the provided encryption key. Note that this is not supported for TPU at the moment.
@@ -83,10 +91,6 @@ def pipeline(
   prompt_column = 'input_text'
   candidate_columns = ['candidate_0', 'candidate_1']
   choice_column = 'choice'
-  machine_spec = function_based.resolve_machine_spec(
-      accelerator_type=accelerator_type,
-      use_test_spec=env.get_use_test_machine_spec(),
-  ).set_display_name('Resolve Machine Spec')
 
   processed_preference_dataset = (
       preprocess_chat_dataset.preprocess_chat_dataset(
@@ -136,16 +140,13 @@ def pipeline(
       .set_caching_options(False)
   )
 
-  reward_model_image_uri = function_based.resolve_private_refined_image_uri(
-      accelerator_type=machine_spec.outputs['accelerator_type'],
-  ).set_display_name('Resolve Reward Model Image URI')
   num_microbatches = function_based.resolve_num_microbatches(
       large_model_reference=reward_model_reference,
   ).set_display_name('Resolve Number of Microbatches')
   reward_model = (
       reward_model_trainer.reward_model_trainer(
           project=project,
-          location=machine_spec.outputs['tuning_location'],
+          location=tuning_location,
           input_model_path=reward_model_path,
           input_dataset_path=preference_dataset_importer.outputs[
               'output_dataset_path'
@@ -154,11 +155,11 @@ def pipeline(
               'output_dataset_path'
           ],
           train_steps=reward_model_train_steps,
-          accelerator_type=machine_spec.outputs['accelerator_type'],
-          accelerator_count=machine_spec.outputs['accelerator_count'],
+          accelerator_type=accelerator_type,
+          accelerator_count=accelerator_count,
           large_model_reference=reward_model_reference,
-          machine_type=machine_spec.outputs['machine_type'],
-          image_uri=reward_model_image_uri.output,
+          machine_type=machine_type,
+          image_uri=reward_model_image_uri,
           inputs_sequence_length=prompt_sequence_length,
           targets_sequence_length=target_sequence_length,
           batch_size=batch_size,

components/google-cloud/google_cloud_pipeline_components/_implementation/llm/rlhf_preprocessor.py

@@ -24,13 +24,25 @@
 @dsl.container_component
 def rlhf_preprocessor(
     large_model_reference: str,
+    accelerator_type: str,
+    use_test_spec: bool,
+    project: str,
+    location: str,
+    artifact_registry: str,
+    tag: str,
     gcp_resources: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
     has_tensorboard_id: dsl.OutputPath(bool),  # pytype: disable=invalid-annotation
     has_inference_dataset: dsl.OutputPath(bool),  # pytype: disable=invalid-annotation
     metadata_large_model_reference: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
     metadata_reference_model_path: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
     metadata_reward_model_reference: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
     metadata_reward_model_path: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    metadata_machine_type: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    metadata_tuning_location: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    metadata_accelerator_type: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    metadata_accelerator_count: dsl.OutputPath(int),  # pytype: disable=invalid-annotation
+    metadata_refined_image_uri: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    use_experimental_image: bool = False,
     evaluation_dataset: str = '',
     tensorboard_resource_id: str = '',
     input_reference_model_path: str = '',
@@ -41,17 +53,30 @@ def rlhf_preprocessor(
 
   Args:
     large_model_reference: The model for fine tuning.
+    accelerator_type: Specific accelerator type for the job.
+    use_test_spec: Whether to use a lower resource machine for testing.
+    project: Project that contains the artifact registry.
+    location: Region that contains the artifact registry.
+    artifact_registry: Registry that contains Docker images.
+    tag: Image tag.
+    use_experimental_image:  Whether to use refined experimental image.
     evaluation_dataset: Path to evaluation data.
     tensorboard_resource_id: TensorBoard resource id.
     metadata_large_model_reference: The base model for fine tuning. The name should be in capitalized snake case format.
     metadata_reference_model_path: The model checkpoint path for the reinforcer model
     metadata_reward_model_reference:  The base model for training reward model. The name should be in capitalized snake case format.
     metadata_reward_model_path: The model checkpoint path for the reward model.
+    image_uri: Docker image URI to use for the custom job.
 
   Returns:
     gcp_resources: GCP resources that can be used to track the custom job.
     has_tensorboard_id: Whether a tensorboard id is provided.
     has_inference_dataset: Whether inference data are provided.
+    metadata_machine_type: The type of the machine to provision for the custom job.
+    metadata_tuning_location: The GCP region to run the custom job.
+    metadata_accelerator_type: Specific accelerator type for the custom job.
+    metadata_accelerator_count: The number of accelerator.
+    metadata_refined_image_uri: Docker image URI to use for the custom job.
   """
   # fmt: on
   return gcpc_utils.build_serverless_customjob_container_spec(
@@ -67,12 +92,24 @@ def rlhf_preprocessor(
               f'--tensorboard_resource_id={tensorboard_resource_id}',
               f'--large_model_reference={large_model_reference}',
               f'--input_reference_model_path={input_reference_model_path}',
+              f'--accelerator_type={accelerator_type}',
+              f'--use_test_spec={use_test_spec}',
+              f'--project={project}',
+              f'--location={location}',
+              f'--artifact_registry={artifact_registry}',
+              f'--tag={tag}',
+              f'--use_experimental_image={use_experimental_image}',
               f'--has_tensorboard_id_path={has_tensorboard_id}',
               f'--has_inference_dataset_path={has_inference_dataset}',
               f'--metadata_large_model_reference_path={metadata_large_model_reference}',
               f'--metadata_reference_model_path_path={metadata_reference_model_path}',
               f'--metadata_reward_model_reference_path={metadata_reward_model_reference}',
               f'--metadata_reward_model_path_path={metadata_reward_model_path}',
+              f'--metadata_machine_type_path={metadata_machine_type}',
+              f'--metadata_tuning_location_path={metadata_tuning_location}',
+              f'--metadata_accelerator_type_path={metadata_accelerator_type}',
+              f'--metadata_accelerator_count_path={metadata_accelerator_count}',
+              f'--metadata_refined_image_uri_path={metadata_refined_image_uri}',
           ],
       ),
       gcp_resources=gcp_resources,

components/google-cloud/google_cloud_pipeline_components/preview/llm/rlhf/component.py

@@ -22,6 +22,7 @@
 from google_cloud_pipeline_components._implementation.llm import reinforcement_learning_graph
 from google_cloud_pipeline_components._implementation.llm import reward_model_graph
 from google_cloud_pipeline_components._implementation.llm import rlhf_preprocessor
+from google_cloud_pipeline_components._implementation.llm import utils
 from google_cloud_pipeline_components._implementation.llm import validate_pipeline
 from google_cloud_pipeline_components.preview.llm.infer import component
 import kfp
@@ -97,6 +98,12 @@ def rlhf_pipeline(
 
   preprocess_metadata = rlhf_preprocessor.rlhf_preprocessor(
       large_model_reference=large_model_reference,
+      accelerator_type=accelerator_type,
+      use_test_spec=env.get_use_test_machine_spec(),
+      project=env.PRIVATE_ARTIFACT_REGISTRY_PROJECT,
+      location=env.PRIVATE_ARTIFACT_REGISTRY_LOCATION,
+      artifact_registry=env.PRIVATE_ARTIFACT_REGISTRY,
+      tag=env.get_private_image_tag(),
       evaluation_dataset=eval_dataset,
       tensorboard_resource_id=tensorboard_resource_id,
   ).set_display_name('Preprocess Inputs')
@@ -112,6 +119,19 @@ def rlhf_pipeline(
               reward_model_path=preprocess_metadata.outputs[
                   'metadata_reward_model_path'
               ],
+              machine_type=preprocess_metadata.outputs['metadata_machine_type'],
+              tuning_location=preprocess_metadata.outputs[
+                  'metadata_tuning_location'
+              ],
+              accelerator_type=preprocess_metadata.outputs[
+                  'metadata_accelerator_type'
+              ],
+              accelerator_count=preprocess_metadata.outputs[
+                  'metadata_accelerator_count'
+              ],
+              reward_model_image_uri=preprocess_metadata.outputs[
+                  'metadata_refined_image_uri'
+              ],
               prompt_sequence_length=prompt_sequence_length,
               target_sequence_length=target_sequence_length,
               eval_dataset=validate_pipeline_task.outputs[
@@ -123,7 +143,6 @@ def rlhf_pipeline(
               lora_dim=reward_lora_dim,
               project=project,
               location=location,
-              accelerator_type=accelerator_type,
               tensorboard_resource_id=tensorboard_resource_id,
               encryption_spec_key_name=encryption_spec_key_name,
           )
@@ -152,6 +171,13 @@ def rlhf_pipeline(
       policy_model_path=preprocess_metadata.outputs[
           'metadata_reference_model_path'
       ],
+      machine_type=preprocess_metadata.outputs['metadata_machine_type'],
+      tuning_location=preprocess_metadata.outputs['metadata_tuning_location'],
+      accelerator_type=preprocess_metadata.outputs['metadata_accelerator_type'],
+      accelerator_count=preprocess_metadata.outputs[
+          'metadata_accelerator_count'
+      ],
+      rl_image_uri=preprocess_metadata.outputs['metadata_refined_image_uri'],
       prompt_sequence_length=prompt_sequence_length,
       target_sequence_length=target_sequence_length,
       reinforcement_learning_rate_multiplier=reinforcement_learning_rate_multiplier,
@@ -160,7 +186,6 @@ def rlhf_pipeline(
       instruction=instruction,
       reward_lora_dim=reward_lora_dim,
       project=project,
-      accelerator_type=accelerator_type,
       location=location,
       tensorboard_resource_id=tensorboard_resource_id,
       encryption_spec_key_name=encryption_spec_key_name,