MLCube integration with Bert

[davidjurado]

MLCube for Bert

MLCube™ GitHub repository. MLCube™ wiki.

Project setup

An important requirement is that you must have Docker installed.

# Create Python environment and install MLCube Docker runner 
virtualenv -p python3 ./env && source ./env/bin/activate && pip install pip==24.0 && pip install mlcube-docker
# Fetch the implementation from GitHub
git clone https://github.com/mlcommons/training && cd ./training/language_model/tensorflow/bert

Go to mlcube directory and study what tasks MLCube implements.

cd ./mlcube
mlcube describe

Demo execution

These tasks will use a demo dataset to execute a faster training workload for a quick demo (~8 min):

mlcube run --task=download_demo -Pdocker.build_strategy=always

mlcube run --task=demo -Pdocker.build_strategy=always

It's also possible to execute the two tasks in one single instruction:

mlcube run --task=download_demo,demo -Pdocker.build_strategy=always

MLCube tasks

Download dataset.

mlcube run --task=download_data -Pdocker.build_strategy=always

Process dataset.

mlcube run --task=process_data -Pdocker.build_strategy=always

Train SSD.

mlcube run --task=train -Pdocker.build_strategy=always

Run compliance checker.

mlcube run --task=check_logs -Pdocker.build_strategy=always

Execute the complete pipeline

You can execute the complete pipeline with one single command.

mlcube run --task=download_data,process_data,train,check_logs -Pdocker.build_strategy=always

TPU Training

For executing this benchmark using TPU you will need access to Google Cloud Platform, then you can create a project (Note: all the resources should be created in the same project) and after that, you will need to follow the next steps:

1. Create a TPU node

In the Google Cloud console, search for the Cloud TPU API page, then click Enable.

Then go to the virtual machine sections and select TPUs

Select create TPU node, fill in all the needed parameters, the recommended TPU type in the readme is v3-128 and the recommended TPU software version is 2.4.0.

The 3 most important parameters you need to remember are: project name, TPU name, and TPU Zone.

After creating, click on the TPU name to see the TPU details, and copy the Service account (should int the format: service-xxxxxxxxxxxx@cloud-tpu.iam.gserviceaccount.com)

2. Create a Google Storage Bucket

Go to Google Storage and create a new Bucket, define the needed parameters.

In the bucket list select the checkbox for the bucket you just created, then click on permissions, after that click on add principal.

In the new principals field paste the Service account from step 1, and then for the roles select, Storage Legacy Bucket Owner, Storage Legacy Bucket Reader and Storage Legacy Bucket Writer. Then click on save, this will allow the TPU to save the checkpoints during training.

3. Create a VM instance

The idea is to create a virtual machine instance containing all the code we will execute using MLCube.

Go to VM instances, then click on create instance and define all the needed parameters (No GPU needed).

IMPORTANT: In the section Identity and API access, check the option Allow full access to all Cloud APIs, this will allow the connection between this VM, the Cloud Storage Bucket and the TPU.

Start the VM, connect to it via SSH, then use this tutorial to install Docker.

After installing Docker, clone the repo and install MLCube and follow the to install MLCube, then go to the path: training/language_model/tensorflow/bert/mlcube

There modify the file at workspace/parameters.yaml and replace it with your data for:

output_gs: your_gs_bucket_name
tpu_name: your_tpu_instance_name
tpu_zone: your_tpu_zone
gcp_project: your_gcp_project

After that run the command:

mlcube run --task=train_tpu --mlcube=mlcube_tpu.yaml -Pdocker.build_strategy=always

This will start the MLCube task that internally in the host VM will send a gRPC with all the data to the TPU through gRPC, then the TPU will get the code to execute and the information of the Cloud Storage Bucket data and will execute the training workload.

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[arjunsuresh]

FROM tensorflow/tensorflow:2.4.0-gpu

The current code is not working with the latest tensorflow (2.12) right? When we tried only the CPU execution is happening (no GPU)

[davidjurado]

Hello @arjunsuresh

Could you please provide more information about the issue you are facing, maybe an screenshot or the error message your are getting, thanks!

[arjunsuresh]

Hi @davidjurado
Actually I was asking if there is a reason why tensorflow2.4 is used here and not the latest one (2.12)?
https://github.com/mlcommons/training/pull/632/files#diff-5c97df48019afb46e0bdfc9099f9dd798a6448ff672b1a1ec1f664500f8409d9

[davidjurado]

Hi @davidjurado Actually I was asking if there is a reason why tensorflow2.4 is used here and not the latest one (2.12)? https://github.com/mlcommons/training/pull/632/files#diff-5c97df48019afb46e0bdfc9099f9dd798a6448ff672b1a1ec1f664500f8409d9

Hello @arjunsuresh,

That's because in the original readme they mentioned that the model was tested using this Tensorflow version, I tested the newer versions and it works as expected, I'm planning to update the Docker base image to match the newest version, thanks.

[nv-rborkar]

@sgpyc as reference owner could you please review this MR so that we can merge it.

[nv-rborkar]

@sgpyc has reviewed & looks ok.