Skip to content

Commit 961dc67

Browse files
rshragarshraga
authored
Add GPU testing config and GPU roberta model tests (#2025)
* added gpu test * separated model tests into cpu and gpu specific versions * removed redundant Test from name * addressed comments and fixed test skip implemenation * fix gpu test train case by sending model input to device * silenced setup script steps, changed skip decorator to unittest, moved gpu tests to separate folder * add back gpu tests * fix lint
1 parent 3e5f77e commit 961dc67

File tree

5 files changed

+130
-13
lines changed

5 files changed

+130
-13
lines changed

‎.github/workflows/test-linux-gpu.yml

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
name: Unit-tests on Linux GPU
2+
3+
on:
4+
pull_request:
5+
push:
6+
branches:
7+
- nightly
8+
- main
9+
- release/*
10+
workflow_dispatch:
11+
12+
env:
13+
CHANNEL: "nightly"
14+
15+
jobs:
16+
tests:
17+
strategy:
18+
matrix:
19+
python_version: ["3.8"]
20+
cuda_arch_version: ["11.6"]
21+
fail-fast: false
22+
uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
23+
with:
24+
runner: linux.g5.4xlarge.nvidia.gpu
25+
repository: pytorch/text
26+
gpu-arch-type: cuda
27+
gpu-arch-version: ${{ matrix.cuda_arch_version }}
28+
timeout: 120
29+
script: |
30+
# Mark Build Directory Safe
31+
git config --global --add safe.directory /__w/text/text
32+
33+
# Set up Environment Variables
34+
export PYTHON_VERSION="${{ matrix.python_version }}"
35+
export VERSION="${{ matrix.cuda_arch_version }}"
36+
export CUDATOOLKIT="pytorch-cuda=${VERSION}"
37+
38+
# Set CHANNEL
39+
if [[ (${GITHUB_EVENT_NAME} = 'pull_request' && (${GITHUB_BASE_REF} = 'release'*)) || (${GITHUB_REF} = 'refs/heads/release'*) ]]; then
40+
export CHANNEL=test
41+
else
42+
export CHANNEL=nightly
43+
fi
44+
45+
# Create Conda Env
46+
conda create --quiet -yp ci_env python="${PYTHON_VERSION}"
47+
conda activate /work/ci_env
48+
python3 -m pip --quiet install cmake>=3.18.0 ninja
49+
conda env update --file ".circleci/unittest/linux/scripts/environment.yml" --prune
50+
51+
# TorchText-specific Setup
52+
printf "* Downloading SpaCy English models\n"
53+
python -m spacy download en_core_web_sm
54+
printf "* Downloading SpaCy German models\n"
55+
python -m spacy download de_core_news_sm
56+
57+
# Install PyTorch and TorchData
58+
set -ex
59+
conda install \
60+
--yes \
61+
--quiet \
62+
-c "pytorch-${CHANNEL}" \
63+
-c nvidia "pytorch-${CHANNEL}"::pytorch[build="*${VERSION}*"] \
64+
"${CUDATOOLKIT}"
65+
printf "Installing torchdata nightly\n"
66+
python3 -m pip install --pre torchdata --extra-index-url https://download.pytorch.org/whl/nightly/cpu --quiet
67+
python3 setup.py develop
68+
python3 -m pip install parameterized --quiet
69+
70+
# Run Tests
71+
python3 -m torch.utils.collect_env
72+
cd test
73+
python3 -m pytest --junitxml=test-results/junit.xml -v --durations 20 torchtext_unittest/models/gpu_tests

‎test/torchtext_unittest/common/case_utils.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import unittest
55
from itertools import zip_longest
66

7+
import torch
78
from torchtext._internal.module_utils import is_module_available
89

910

@@ -37,6 +38,17 @@ def get_temp_path(self, *paths):
3738
return path
3839

3940

41+
class TestBaseMixin:
42+
"""Mixin to provide consistent way to define device/dtype/backend aware TestCase"""
43+
44+
dtype = None
45+
device = None
46+
47+
def setUp(self):
48+
super().setUp()
49+
torch.random.manual_seed(2434)
50+
51+
4052
def skipIfNoModule(module, display_name=None):
4153
display_name = display_name or module
4254
return unittest.skipIf(not is_module_available(module), f'"{display_name}" is not available')

‎test/torchtext_unittest/models/gpu_tests/models_gpu_test.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
import unittest
2+
3+
import torch
4+
from torchtext_unittest.common.torchtext_test_case import TorchtextTestCase
5+
from torchtext_unittest.models.models_test_impl import BaseTestModels
6+
7+
8+
@unittest.skipIf(not torch.cuda.is_available(), reason="CUDA is not available")
9+
class TestModels32GPU(BaseTestModels, TorchtextTestCase):
10+
dtype = torch.float32
11+
device = torch.device("cuda")

‎test/torchtext_unittest/models/models_cpu_test.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
import torch
2+
3+
from ..common.torchtext_test_case import TorchtextTestCase
4+
from .models_test_impl import BaseTestModels
5+
6+
7+
class TestModels32CPU(BaseTestModels, TorchtextTestCase):
8+
dtype = torch.float32
9+
device = torch.device("cpu")

‎test/torchtext_unittest/models/test_models.py renamed to ‎test/torchtext_unittest/models/models_test_impl.py

Lines changed: 25 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -2,30 +2,42 @@
22
from unittest.mock import patch
33

44
import torch
5-
import torchtext
65
from torch.nn import functional as torch_F
76

8-
from ..common.torchtext_test_case import TorchtextTestCase
7+
from ..common.case_utils import TestBaseMixin
98

109

11-
class TestModels(TorchtextTestCase):
10+
class BaseTestModels(TestBaseMixin):
11+
def get_model(self, encoder_conf, head=None, freeze_encoder=False, checkpoint=None, override_checkpoint_head=False):
12+
from torchtext.models import RobertaBundle
13+
14+
model = RobertaBundle.build_model(
15+
encoder_conf=encoder_conf,
16+
head=head,
17+
freeze_encoder=freeze_encoder,
18+
checkpoint=checkpoint,
19+
override_checkpoint_head=override_checkpoint_head,
20+
)
21+
model.to(device=self.device, dtype=self.dtype)
22+
return model
23+
1224
def test_roberta_bundler_build_model(self) -> None:
13-
from torchtext.models import RobertaClassificationHead, RobertaEncoderConf, RobertaModel, RobertaBundle
25+
from torchtext.models import RobertaClassificationHead, RobertaEncoderConf, RobertaModel
1426

1527
dummy_encoder_conf = RobertaEncoderConf(
1628
vocab_size=10, embedding_dim=16, ffn_dimension=64, num_attention_heads=2, num_encoder_layers=2
1729
)
1830

1931
# case: user provide encoder checkpoint state dict
2032
dummy_encoder = RobertaModel(dummy_encoder_conf)
21-
model = RobertaBundle.build_model(encoder_conf=dummy_encoder_conf, checkpoint=dummy_encoder.state_dict())
33+
model = self.get_model(encoder_conf=dummy_encoder_conf, checkpoint=dummy_encoder.state_dict())
2234
self.assertEqual(model.state_dict(), dummy_encoder.state_dict())
2335

2436
# case: user provide classifier checkpoint state dict when head is given and override_head is False (by default)
2537
dummy_classifier_head = RobertaClassificationHead(num_classes=2, input_dim=16)
2638
another_dummy_classifier_head = RobertaClassificationHead(num_classes=2, input_dim=16)
2739
dummy_classifier = RobertaModel(dummy_encoder_conf, dummy_classifier_head)
28-
model = RobertaBundle.build_model(
40+
model = self.get_model(
2941
encoder_conf=dummy_encoder_conf,
3042
head=another_dummy_classifier_head,
3143
checkpoint=dummy_classifier.state_dict(),
@@ -34,7 +46,7 @@ def test_roberta_bundler_build_model(self) -> None:
3446

3547
# case: user provide classifier checkpoint state dict when head is given and override_head is set True
3648
another_dummy_classifier_head = RobertaClassificationHead(num_classes=2, input_dim=16)
37-
model = RobertaBundle.build_model(
49+
model = self.get_model(
3850
encoder_conf=dummy_encoder_conf,
3951
head=another_dummy_classifier_head,
4052
checkpoint=dummy_classifier.state_dict(),
@@ -48,13 +60,13 @@ def test_roberta_bundler_build_model(self) -> None:
4860
encoder_state_dict = {}
4961
for k, v in dummy_classifier.encoder.state_dict().items():
5062
encoder_state_dict["encoder." + k] = v
51-
model = torchtext.models.RobertaBundle.build_model(
63+
model = self.get_model(
5264
encoder_conf=dummy_encoder_conf, head=dummy_classifier_head, checkpoint=encoder_state_dict
5365
)
5466
self.assertEqual(model.state_dict(), dummy_classifier.state_dict())
5567

5668
def test_roberta_bundler_train(self) -> None:
57-
from torchtext.models import RobertaClassificationHead, RobertaEncoderConf, RobertaModel, RobertaBundle
69+
from torchtext.models import RobertaClassificationHead, RobertaEncoderConf, RobertaModel
5870

5971
dummy_encoder_conf = RobertaEncoderConf(
6072
vocab_size=10, embedding_dim=16, ffn_dimension=64, num_attention_heads=2, num_encoder_layers=2
@@ -63,8 +75,8 @@ def test_roberta_bundler_train(self) -> None:
6375

6476
def _train(model):
6577
optim = SGD(model.parameters(), lr=1)
66-
model_input = torch.tensor([[0, 1, 2, 3, 4, 5]])
67-
target = torch.tensor([0])
78+
model_input = torch.tensor([[0, 1, 2, 3, 4, 5]]).to(device=self.device)
79+
target = torch.tensor([0]).to(device=self.device)
6880
logits = model(model_input)
6981
loss = torch_F.cross_entropy(logits, target)
7082
loss.backward()
@@ -73,7 +85,7 @@ def _train(model):
7385
# does not freeze encoder
7486
dummy_classifier_head = RobertaClassificationHead(num_classes=2, input_dim=16)
7587
dummy_classifier = RobertaModel(dummy_encoder_conf, dummy_classifier_head)
76-
model = RobertaBundle.build_model(
88+
model = self.get_model(
7789
encoder_conf=dummy_encoder_conf,
7890
head=dummy_classifier_head,
7991
freeze_encoder=False,
@@ -91,7 +103,7 @@ def _train(model):
91103
# freeze encoder
92104
dummy_classifier_head = RobertaClassificationHead(num_classes=2, input_dim=16)
93105
dummy_classifier = RobertaModel(dummy_encoder_conf, dummy_classifier_head)
94-
model = RobertaBundle.build_model(
106+
model = self.get_model(
95107
encoder_conf=dummy_encoder_conf,
96108
head=dummy_classifier_head,
97109
freeze_encoder=True,

0 commit comments

Comments
 (0)