Use pytorch 2 image

determined-ai · Feb 26, 2024 · 353939e · 353939e
1 parent 31c8dfb
commit 353939e
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 2 deletions.
diff --git a/blog/llm-finetuning-2/README.md b/blog/llm-finetuning-2/README.md
@@ -32,7 +32,7 @@ DeepSpeed configuration files are in the `ds_configs` folder.
 Test your model's generation capabilities:
 
 ```bash
-python test_model.py --exp_id <exp_id> --dataset_subset <dataset_subset>
+python inference.py --exp_id <exp_id> --dataset_subset <dataset_subset>
 ```
 
 Where 
@@ -42,7 +42,7 @@ Where
 To test the pretrained model (not finetuned), leave out `--exp_id`. For example:
 
 ```bash
-python test_model.py --dataset_subset easy
+python inference.py --dataset_subset easy
 ```
 
 ## Validating the tokenizer

diff --git a/blog/llm-finetuning-2/deepspeed.yaml b/blog/llm-finetuning-2/deepspeed.yaml
@@ -3,6 +3,7 @@ debug: false
 environment:
   environment_variables:
     - NCCL_DEBUG=INFO
+  image: determinedai/environments:cuda-11.8-pytorch-2.0-gpu-mpi-0.27.1
 resources:
   slots_per_trial: 2
 searcher:

diff --git a/blog/llm-finetuning-2/lora.yaml b/blog/llm-finetuning-2/lora.yaml
@@ -3,6 +3,7 @@ debug: false
 environment:
   environment_variables:
     - NCCL_DEBUG=INFO
+  image: determinedai/environments:cuda-11.8-pytorch-2.0-gpu-mpi-0.27.1
 resources:
   slots_per_trial: 2
 searcher: