From 95b3e2b3a39580177ef4ce8c6ee98a36532828fb Mon Sep 17 00:00:00 2001 From: deepanker13 Date: Fri, 12 Jan 2024 01:27:12 +0530 Subject: [PATCH] gpu training example fix --- examples/sdk/train_api.ipynb | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/examples/sdk/train_api.ipynb b/examples/sdk/train_api.ipynb index 76c74a0354..242efe8ebb 100644 --- a/examples/sdk/train_api.ipynb +++ b/examples/sdk/train_api.ipynb @@ -73,9 +73,12 @@ " ),\n", " training_parameters=TrainingArguments(\n", " num_train_epochs=1,\n", - " per_device_train_batch_size=4,\n", - " gradient_accumulation_steps=4,\n", + " per_device_train_batch_size=1,\n", + " gradient_accumulation_steps=1,\n", " gradient_checkpointing=True,\n", + " gradient_checkpointing_kwargs={\n", + " \"use_reentrant\": False\n", + " }, # this is mandatory if checkpointng is enabled\n", " warmup_steps=0.02,\n", " learning_rate=1,\n", " lr_scheduler_type=\"cosine\",\n", @@ -93,7 +96,7 @@ " resources_per_worker={\n", " \"gpu\": 1,\n", " \"cpu\": 8,\n", - " \"memory\": \"16Gi\",\n", + " \"memory\": \"8Gi\",\n", " }, # remove the gpu key if you don't want to attach gpus to the pods\n", ")" ]