[wip] makes it possible to run alpaca with flyte

Dockerfile

@@ -0,0 +1,22 @@
+FROM pytorch/pytorch:2.0.0-cuda11.7-cudnn8-runtime
+
+LABEL org.opencontainers.image.source https://github.com/unionai-oss/stanford-alpaca
+
+WORKDIR /root
+ENV VENV /opt/venv
+ENV LANG C.UTF-8
+ENV LC_ALL C.UTF-8
+ENV PYTHONPATH /root
+
+ARG VERSION
+ARG DOCKER_IMAGE
+
+RUN apt-get update && apt-get install build-essential -y
+
+COPY . /root
+
+WORKDIR /root
+# Pod tasks should be exposed in the default image
+RUN pip install -r requirements.txt
+
+ENV FLYTE_INTERNAL_IMAGE "$DOCKER_IMAGE"

requirements.txt

@@ -7,3 +7,6 @@ torch
 sentencepiece
 tokenizers==0.12.1
 wandb
+https://github.com/flyteorg/flyteidl/archive/e435d9146a36c94e8ee8acd193063485c2da89aa.zip#egg=flyteidl
+https://github.com/flyteorg/flytekit/archive/e97780b25b6539b302876bf68e5a94f96d3fb9b8.zip#egg=flytekit
+https://github.com/flyteorg/flytekit/archive/e97780b25b6539b302876bf68e5a94f96d3fb9b8.zip#subdirectory=plugins/flytekit-kf-pytorch&egg=flytekitplugins-kfpytorch

train.py

@@ -14,15 +14,20 @@
 
 import copy
 import logging
+import os
 from dataclasses import dataclass, field
 from typing import Optional, Dict, Sequence
 
 import torch
 import transformers
+from flytekit import Resources
 from torch.utils.data import Dataset
 from transformers import Trainer
 
 import utils
+import flytekit
+from flytekitplugins.kfpytorch.task import Elastic
+from dataclasses_json import dataclass_json
 
 IGNORE_INDEX = -100
 DEFAULT_PAD_TOKEN = "[PAD]"
@@ -43,16 +48,19 @@
 }
 
 
+@dataclass_json
 @dataclass
 class ModelArguments:
     model_name_or_path: Optional[str] = field(default="facebook/opt-125m")
 
 
+@dataclass_json
 @dataclass
 class DataArguments:
     data_path: str = field(default=None, metadata={"help": "Path to the training data."})
 
 
+@dataclass_json
 @dataclass
 class TrainingArguments(transformers.TrainingArguments):
     cache_dir: Optional[str] = field(default=None)
@@ -189,10 +197,18 @@ def make_supervised_data_module(tokenizer: transformers.PreTrainedTokenizer, dat
     return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator)
 
 
-def train():
-    parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
-    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
-
+# TODO: update the Dockerfile for use with cuda
+@flytekit.task(
+    task_config=Elastic(nnodes=1),
+    environment={
+        "TRANSFORMERS_CACHE": "/tmp",
+        "WANDB_API_KEY": "<wandb_key>",
+        "WANDB_PROJECT": "unioncloud-llms",
+    },
+    requests=Resources(mem="40Gi", cpu="1", gpu="1"),
+)
+def train(model_args: ModelArguments, data_args: DataArguments, training_args: TrainingArguments) -> flytekit.file.FlyteFile:
+    os.environ["WANDB_RUN_ID"] = os.environ.get("FLYTE_INTERNAL_EXECUTION_ID", "foo")
     model = transformers.AutoModelForCausalLM.from_pretrained(
         model_args.model_name_or_path,
         cache_dir=training_args.cache_dir,
@@ -225,6 +241,18 @@ def train():
     trainer.train()
     trainer.save_state()
     safe_save_model_for_hf_trainer(trainer=trainer, output_dir=training_args.output_dir)
+    return flytekit.file.FlyteFile(path=training_args.output_dir)
+
+
+@flytekit.workflow
+def train_wf(model_args: ModelArguments, data_args: DataArguments, training_args: TrainingArguments) -> flytekit.file.FlyteFile:
+    return train(model_args=model_args, data_args=data_args, training_args=training_args)
+
+
+def train_cli():
+    parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
+    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+    train(model_args, data_args, training_args)
 
 
 if __name__ == "__main__":