Skip to content

Commit

Permalink
feat: LLM - Added new regions for tuning and tuned model inference
Browse files Browse the repository at this point in the history 
PiperOrigin-RevId: 576294181
  • Loading branch information
@Ark-kun @copybara-github
Ark-kun authored and copybara-github committed Oct 24, 2023
1 parent ffbd872 commit 3d43497
Showing 1 changed file with 37 additions and 20 deletions.
57 changes: 37 additions & 20 deletions vertexai/language_models/_language_models.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -190,8 +190,7 @@ def tune_model(
Learning rate to use in tuning.
learning_rate_multiplier: Learning rate multiplier to use in tuning.
tuning_job_location: GCP location where the tuning job should be run.
Only "europe-west4" and "us-central1" locations are supported for now.
tuned_model_location: GCP location where the tuned model should be deployed. Only "us-central1" is supported for now.
tuned_model_location: GCP location where the tuned model should be deployed.
model_display_name: Custom display name for the tuned model.
tuning_evaluation_spec: Specification for the model evaluation during tuning.
default_context: The context to use for all training samples by default.
Expand Down Expand Up @@ -294,8 +293,7 @@ def _tune_model(
See https://cloud.google.com/vertex-ai/docs/generative-ai/models/tune-models#dataset_format
tuning_parameters: Tuning pipeline parameter values.
tuning_job_location: GCP location where the tuning job should be run.
Only "europe-west4" and "us-central1" locations are supported for now.
tuned_model_location: GCP location where the tuned model should be deployed. Only "us-central1" is supported for now.
tuned_model_location: GCP location where the tuned model should be deployed.
model_display_name: Custom display name for the tuned model.
Returns:
Expand All @@ -312,9 +310,10 @@ def _tune_model(
"Please specify the tuning job location (`tuning_job_location`)."
f"Tuning is supported in the following locations: {_TUNING_LOCATIONS}"
)
if tuned_model_location != _TUNED_MODEL_LOCATION:
if tuned_model_location not in _TUNED_MODEL_LOCATIONS:
raise ValueError(
f'Model deployment is only supported in the following locations: tuned_model_location="{_TUNED_MODEL_LOCATION}"'
"Tuned model deployment is only supported in the following locations: "
f"{_TUNED_MODEL_LOCATIONS}"
)
model_info = _model_garden_models._get_model_info(
model_id=self._model_id,
Expand All @@ -329,6 +328,7 @@ def _tune_model(
tuning_parameters=tuning_parameters,
model_display_name=model_display_name,
tuning_job_location=tuning_job_location,
tuned_model_location=tuned_model_location,
)

job = _LanguageModelTuningJob(
Expand Down Expand Up @@ -369,8 +369,7 @@ def tune_model(
train_steps: Number of training batches to tune on (batch size is 8 samples).
learning_rate_multiplier: Learning rate multiplier to use in tuning.
tuning_job_location: GCP location where the tuning job should be run.
Only "europe-west4" and "us-central1" locations are supported for now.
tuned_model_location: GCP location where the tuned model should be deployed. Only "us-central1" is supported for now.
tuned_model_location: GCP location where the tuned model should be deployed.
model_display_name: Custom display name for the tuned model.
tuning_evaluation_spec: Specification for the model evaluation during tuning.
accelerator_type: Type of accelerator to use. Can be "TPU" or "GPU".
Expand Down Expand Up @@ -436,8 +435,7 @@ def tune_model(
Learning rate to use in tuning.
learning_rate_multiplier: Learning rate multiplier to use in tuning.
tuning_job_location: GCP location where the tuning job should be run.
Only "europe-west4" and "us-central1" locations are supported for now.
tuned_model_location: GCP location where the tuned model should be deployed. Only "us-central1" is supported for now.
tuned_model_location: GCP location where the tuned model should be deployed.
model_display_name: Custom display name for the tuned model.
tuning_evaluation_spec: Specification for the model evaluation during tuning.
accelerator_type: Type of accelerator to use. Can be "TPU" or "GPU".
Expand Down Expand Up @@ -503,8 +501,7 @@ def tune_model(
Learning rate to use in tuning.
learning_rate_multiplier: Learning rate multiplier to use in tuning.
tuning_job_location: GCP location where the tuning job should be run.
Only "europe-west4" and "us-central1" locations are supported for now.
tuned_model_location: GCP location where the tuned model should be deployed. Only "us-central1" is supported for now.
tuned_model_location: GCP location where the tuned model should be deployed.
model_display_name: Custom display name for the tuned model.
default_context: The context to use for all training samples by default.
accelerator_type: Type of accelerator to use. Can be "TPU" or "GPU".
Expand Down Expand Up @@ -570,8 +567,7 @@ def tune_model(
Learning rate to use in tuning.
learning_rate_multiplier: Learning rate multiplier to use in tuning.
tuning_job_location: GCP location where the tuning job should be run.
Only "europe-west4" and "us-central1" locations are supported for now.
tuned_model_location: GCP location where the tuned model should be deployed. Only "us-central1" is supported for now.
tuned_model_location: GCP location where the tuned model should be deployed.
model_display_name: Custom display name for the tuned model.
default_context: The context to use for all training samples by default.
accelerator_type: Type of accelerator to use. Can be "TPU" or "GPU".
Expand Down Expand Up @@ -2697,9 +2693,29 @@ class _PreviewCodeGenerationModel(

###### Model tuning
# Currently, tuning can only work in this location
_TUNING_LOCATIONS = ("europe-west4", "us-central1")
# Currently, deployment can only work in this location
_TUNED_MODEL_LOCATION = "us-central1"

_SUPPORTED_LOCATIONS = [
# 1
"us-central1",
"europe-west4",
"asia-southeast1",
# 2
"us-west1",
"europe-west3",
"europe-west2",
"asia-northeast1",
# 3
"us-east4",
"us-west4",
"northamerica-northeast1",
"europe-west9",
"europe-west1",
"asia-northeast3",
]

_TUNING_LOCATIONS = _SUPPORTED_LOCATIONS
# Currently, deployment can only work in these locations
_TUNED_MODEL_LOCATIONS = _SUPPORTED_LOCATIONS


class _LanguageModelTuningJob:
Expand Down Expand Up @@ -2771,7 +2787,7 @@ def _list_tuned_model_names(model_id: str) -> List[str]:
tuned_models = aiplatform.Model.list(
filter=f'labels.{_TUNING_BASE_MODEL_ID_LABEL_KEY}="{model_id.replace("@", "-")}"',
# TODO(b/275444096): Remove the explicit location once models are deployed to the user's selected location
location=_TUNED_MODEL_LOCATION,
location=aiplatform_initializer.global_config.location,
)
model_names = [model.resource_name for model in tuned_models]
return model_names
Expand All @@ -2789,8 +2805,9 @@ def _launch_tuning_job(
model_id: str,
tuning_pipeline_uri: str,
tuning_parameters: Dict[str, Any],
tuning_job_location: str,
tuned_model_location: str,
model_display_name: Optional[str] = None,
tuning_job_location: str = _TUNING_LOCATIONS[0],
) -> aiplatform.PipelineJob:
output_dir_uri = _generate_tuned_model_dir_uri(model_id=model_id)
if isinstance(training_data, str):
Expand Down Expand Up @@ -2836,7 +2853,7 @@ def _launch_tuning_job(
"project": aiplatform_initializer.global_config.project,
# TODO(b/275444096): Remove the explicit location once tuning can happen in all regions
# "location": aiplatform_initializer.global_config.location,
"location": _TUNED_MODEL_LOCATION,
"location": tuned_model_location,
"large_model_reference": model_id,
"model_display_name": model_display_name,
}
Expand Down

0 comments on commit 3d43497

Please sign in to comment.