Skip to content

Commit

Permalink
Remove MonkeyPatch for GPUs (#1227)
Browse files Browse the repository at this point in the history
SUMMARY:
- This was added previously due to an error we were seeing with multiple
GPUs (i.e having more than one gpu visible would cause an error)
- Can verify this no longer happens
  • Loading branch information
dsikka authored Mar 5, 2025
1 parent 7bb517f commit 07726ef
Show file tree
Hide file tree
Showing 2 changed files with 0 additions and 11 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,6 @@ def create_mock_file(self, extension, content, path, filename):

def tearDown(self):
shutil.rmtree(self.output)
self.monkeypatch.undo()


@pytest.mark.integration
Expand All @@ -121,11 +120,8 @@ class TestOneshotCustomDatasetSmall(TestFinetuneNoRecipeCustomDataset):
def setUp(self):
import torch

self.monkeypatch = pytest.MonkeyPatch()

if torch.cuda.is_available():
self.device = "cuda:0"
self.monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0")
else:
self.device = "cpu"

Expand All @@ -147,15 +143,12 @@ def setUp(self):
import torch
from transformers import AutoModelForCausalLM

self.monkeypatch = pytest.MonkeyPatch()
self.device = "cuda:0"
self.output = "./oneshot_output"
self.monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0")

self.model = AutoModelForCausalLM.from_pretrained(
self.model, device_map=self.device, torch_dtype=torch.bfloat16
)
self.monkeypatch = pytest.MonkeyPatch()

def test_oneshot_then_finetune_gpu(self):
self._test_finetune_wout_recipe_custom_dataset()
Original file line number Diff line number Diff line change
Expand Up @@ -21,16 +21,13 @@ def setUp(self):
self.output = "./finetune_output"
# finetune workflows in general seem to have trouble with multi-gpus
# use just one atm
self.monkeypatch = pytest.MonkeyPatch()

def test_oneshot_and_finetune_with_tokenizer(self):
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer

from llmcompressor.transformers import compress

self.monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0")

recipe_str = (
"tests/llmcompressor/transformers/finetune/test_alternate_recipe.yaml"
)
Expand Down Expand Up @@ -71,4 +68,3 @@ def test_oneshot_and_finetune_with_tokenizer(self):

def tearDown(self):
shutil.rmtree(self.output)
self.monkeypatch.undo()

0 comments on commit 07726ef

Please sign in to comment.