How to Load model.tar.gz from S3 for Fine-Tuning with SageMaker Training Jobs
Tags: aws, deeplearning, beginners, python
Published: 2023-05-03 06:35:56

Having previously conducted pre-training for a Tabular BERT model with SageMaker Training Jobs, I proceeded to create a job for fine-tuning using the `model.tar.gz` file from the pre-training results.

Although the job is fundamentally similar to the pre-training process, certain adjustments were required in the following aspects, which I have compiled as a memo:

- Extracting the tar file
- Switching arguments between local environment and SageMaker using environment variables

## Extracting the tar file
The pre-training job has saved the `model.tar.gz` file on S3.
![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/yl434ffb0hxgl6jxlve3.png)

The `model.tar.gz` file, which is stored on S3 as a result of the pre-training job, contains the following files: the model `pytorch_model.bin`, the configuration file `config.json`, the dictionary file `vocab.nb`, and the token-to-id conversion file `vocab_token2id.bin`. To load these files during fine-tuning, it is essential to devise a method for extracting the tar file upon job execution.

Initially, set the S3 path for the `model.tar.gz` file in the input_model section of the job file. Consequently, the `model.tar.gz` file will be placed in the `/opt/ml/input/data/input_model/` (model_path) directory when the job is executed.

```python
import sagemaker
from sagemaker.estimator import Estimator

session = sagemaker.Session()
role = sagemaker.get_execution_role()

estimator = Estimator(
    image_uri=<image-url>,
    role=role,
    instance_type="ml.g4dn.2xlarge",
    instance_count=1,
    base_job_name="tabformer-opt-fine-tuning",
    output_path="s3://<bucket-name>/sagemaker/output_data/fine_tuning",
    code_location="s3://<bucket-name>/sagemaker/output_data/fine_tuning",
    sagemaker_session=session,
    entry_point="fine-tuning.sh",
    dependencies=["tabformer-opt"],
    hyperparameters={
        "data_root": "/opt/ml/input/data/input_data/",
        "data_fname": "summary",
        "output_dir": "/opt/ml/model/",
        "model_path": "/opt/ml/input/data/input_model/",
    }
)
estimator.fit({
    "input_data": "s3://<bucket-name>/sagemaker/input_data/summary.csv",
    "input_model": "s3://<bucket-name>/sagemaker/output_data/pre_training/tabformer-opt-2022-12-16-07-00-45-931/output/model.tar.gz"
})
```
Next, include the following in the fine-tuning execution file `tabformer_bert_fine_tuning.py`:

```python
        with tarfile.open(name=path.join(args.model_path, f'model.tar.gz'), mode="r:gz") as mytar:
            mytar.extractall(path.join(args.model_path, f'model'))
            
            token2id_file = path.join(args.model_path, f"model/vocab_token2id.bin")
            vocab_file = path.join(args.model_path, f"model/vocab.nb")
            pretrained_model = path.join(args.model_path, f"model/checkpoint-500/pytorch_model.bin")
            pretrained_config = path.join(args.model_path, f"model/checkpoint-500/config.json")
```
The `tarfile.open()` function reads the `model.tar.gz` file, and `mytar.extractall(path.join(args.model_path, f'model'))` extracts the contents under the `/opt/ml/input/data/input_model/model/` directory.

This allows you to load the extracted files, such as with `token2id_file = path.join(args.model_path, f"model/vocab_token2id.bin")`.

## Switching arguments between local environment and SageMaker using environment variables
With this setup, you can now load the `model.tar.gz` file from S3. However, there may be cases where you want to change the source of the file when performing fine-tuning locally.

To handle such situations, you can use `os.getenv('SM_MODEL_DIR')` to obtain the SageMaker environment variable `SM_MODEL_DIR` (the directory that will be uploaded to S3 upon container termination) and switch the source of the file between local and SageMaker (Job) environments.

```python
    key = os.getenv('SM_MODEL_DIR')
    
    if key :
        with tarfile.open(name=path.join(args.model_path, f'model.tar.gz'), mode="r:gz") as mytar:
            mytar.extractall(path.join(args.model_path, f'model'))
            
            token2id_file = path.join(args.model_path, f"model/vocab_token2id.bin")
            vocab_file = path.join(args.model_path, f"model/vocab.nb")
            pretrained_model = path.join(args.model_path, f"model/checkpoint-500/pytorch_model.bin")
            pretrained_config = path.join(args.model_path, f"model/checkpoint-500/config.json")
    else :
            vocab_file = path.join(args.model_path, f"vocab.nb")
            token2id_file = path.join(args.model_path, f"vocab_token2id.bin")
            pretrained_model = path.join(args.model_path, f"checkpoint-500/pytorch_model.bin")
            pretrained_config = path.join(args.model_path, f"checkpoint-500/config.json")
```