deepspeed

1
# Copyright (c) Microsoft Corporation.
2
# SPDX-License-Identifier: Apache-2.0
3

4
# DeepSpeed Team
5
"""
6
Note: please copy webtext data to "Megatron-LM" folder, before running this script.
7
"""
8

9
import sys
10
import unittest
11

12
sys.path.append('../DeepSpeedExamples/Megatron_GPT2')
13
sys.path.append('../DeepSpeedExamples/BingBertSquad')
14

15
# Import the test cases here.
16
import Megatron_GPT2
17
import BingBertSquad
18

19

20
def pytest_hack(runner_result):
21
    '''This is an ugly hack to get the unittest suites to play nicely with
22
    pytest. Otherwise failed tests are not reported by pytest for some reason.
23

24
    Long-term, these model tests should be adapted to pytest.
25
    '''
26
    if not runner_result.wasSuccessful():
27
        print('SUITE UNSUCCESSFUL:', file=sys.stderr)
28
        for fails in runner_result.failures:
29
            print(fails, file=sys.stderr)
30
        assert runner_result.wasSuccessful()  # fail the test
31

32

33
def test_megatron():
34
    runner = unittest.TextTestRunner(failfast=True)
35
    pytest_hack(runner.run(Megatron_GPT2.suite()))
36

37

38
def test_megatron_checkpoint():
39
    runner = unittest.TextTestRunner(failfast=True)
40
    pytest_hack(runner.run(Megatron_GPT2.checkpoint_suite()))
41

42

43
def test_squad():
44
    runner = unittest.TextTestRunner(failfast=True)
45
    pytest_hack(runner.run(BingBertSquad.suite()))
46