"examples/model_compress/models/cifar10/vgg.py" did not exist on "5882039607b92904f2c5dd124d58e26b0dc8ab92"
test_deepspeed.py 48.1 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
# Copyright 2020 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

15
import dataclasses
16
import io
17
import itertools
18
import json
19
20
import os
import unittest
21
from copy import deepcopy
22
from functools import partial
23

24
import datasets
25
from parameterized import parameterized
26

27
import tests.trainer.test_trainer
Stas Bekman's avatar
Stas Bekman committed
28
from tests.trainer.test_trainer import TrainerIntegrationCommon  # noqa
29
from transformers import AutoModel, TrainingArguments, is_torch_available, logging
30
31
32
33
34
from transformers.integrations.deepspeed import (
    HfDeepSpeedConfig,
    is_deepspeed_available,
    unset_hf_deepspeed_config,
)
35
from transformers.testing_utils import (
36
    CaptureLogger,
37
    CaptureStd,
38
    CaptureStderr,
39
    LoggingLevel,
40
41
42
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
43
    mockenv_context,
44
    require_deepspeed,
45
    require_optuna,
46
47
48
49
    require_torch_gpu,
    require_torch_multi_gpu,
    slow,
)
50
from transformers.trainer_utils import get_last_checkpoint, set_seed
51
from transformers.utils import WEIGHTS_NAME, is_torch_bf16_gpu_available
52

53

54
if is_torch_available():
Stas Bekman's avatar
Stas Bekman committed
55
56
57
58
    from tests.trainer.test_trainer import (  # noqa
        RegressionModelConfig,
        RegressionPreTrainedModel,
    )
59

60
61
62
    # hack to restore original logging level pre #21700
    get_regression_trainer = partial(tests.trainer.test_trainer.get_regression_trainer, log_level="info")

63

64
set_seed(42)
65

66
67
68
# default torch.distributed port
DEFAULT_MASTER_PORT = "10999"

69
T5_SMALL = "t5-small"
70
T5_TINY = "patrickvonplaten/t5-tiny-random"
71
GPT2_TINY = "sshleifer/tiny-gpt2"
72
73


74
75
76
77
78
def load_json(path):
    with open(path) as f:
        return json.load(f)


Stas Bekman's avatar
Stas Bekman committed
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
def get_master_port(real_launcher=False):
    """
    When using a single gpu launcher emulation (i.e. not deepspeed or python -m torch.distributed)
    the issue is that once the port is tied it can't be used anywhere else outside of this process,
    since torch.dist doesn't free the port until the process exits. Therefore for the sake of being
    able to run both emulated launcher and normal launcher tests we need 2 distinct ports.

    This function will give the right port in the right context. For real launcher it'll give the
    base port, for emulated launcher it'll give the base port + 1. In both cases a string is
    returned.

    Args:
        `real_launcher`: whether a real launcher is going to be used, or the emulated one

    """

    master_port_base = os.environ.get("DS_TEST_PORT", DEFAULT_MASTER_PORT)
    if not real_launcher:
        master_port_base = str(int(master_port_base) + 1)
    return master_port_base


101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
def require_deepspeed_aio(test_case):
    """
    Decorator marking a test that requires deepspeed aio (nvme)
    """
    if not is_deepspeed_available():
        return unittest.skip("test requires deepspeed")(test_case)

    import deepspeed
    from deepspeed.ops.aio import AsyncIOBuilder

    if not deepspeed.ops.__compatible_ops__[AsyncIOBuilder.NAME]:
        return unittest.skip("test requires deepspeed async-io")(test_case)
    else:
        return test_case


117
118
if is_deepspeed_available():
    from deepspeed.utils import logger as deepspeed_logger  # noqa
119
    from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
120
    from transformers.integrations.deepspeed import deepspeed_config, is_deepspeed_zero3_enabled  # noqa
121

122
123
124
125
126
127
128

def get_launcher(distributed=False):
    # 1. explicitly set --num_nodes=1 just in case these tests end up run on a multi-node setup
    # - it won't be able to handle that
    # 2. for now testing with just 2 gpus max (since some quality tests may give different
    # results with mode gpus because we use very little data)
    num_gpus = min(2, get_gpu_count()) if distributed else 1
Stas Bekman's avatar
Stas Bekman committed
129
    master_port = get_master_port(real_launcher=True)
130
    return f"deepspeed --num_nodes 1 --num_gpus {num_gpus} --master_port {master_port}".split()
131
132


133
134
ZERO2 = "zero2"
ZERO3 = "zero3"
135
136
137
138

FP16 = "fp16"
BF16 = "bf16"

139
140
141
142
143
144
145
146
HF_OPTIM = "hf_optim"
HF_SCHEDULER = "hf_scheduler"
DS_OPTIM = "ds_optim"
DS_SCHEDULER = "ds_scheduler"

optims = [HF_OPTIM, DS_OPTIM]
schedulers = [HF_SCHEDULER, DS_SCHEDULER]

147
stages = [ZERO2, ZERO3]
148
if is_torch_bf16_gpu_available():
149
150
151
152
153
154
155
156
157
158
159
160
161
162
    dtypes = [FP16, BF16]
else:
    dtypes = [FP16]


def parameterized_custom_name_func(func, param_num, param):
    # customize the test name generator function as we want both params to appear in the sub-test
    # name, as by default it shows only the first param
    param_based_name = parameterized.to_safe_name("_".join(str(x) for x in param.args))
    return f"{func.__name__}_{param_based_name}"


# Cartesian-product of zero stages with models to test
params = list(itertools.product(stages, dtypes))
163

164
165
params_with_optims_and_schedulers = list(itertools.product(stages, dtypes, optims, schedulers))

166

167
168
169
170
171
172
173
174
175
176
@require_deepspeed
@require_torch_gpu
class CoreIntegrationDeepSpeed(TestCasePlus, TrainerIntegrationCommon):
    """
    Testing non-Trainer DeepSpeed integration
    """

    def setUp(self):
        super().setUp()

Stas Bekman's avatar
Stas Bekman committed
177
        master_port = get_master_port(real_launcher=False)
178
179
180
181
182
183
184
        self.dist_env_1_gpu = {
            "MASTER_ADDR": "localhost",
            "MASTER_PORT": master_port,
            "RANK": "0",
            "LOCAL_RANK": "0",
            "WORLD_SIZE": "1",
        }
185

186
187
188
189
190
191
    def tearDown(self):
        super().tearDown()

        # reset the ds config global so that tests state doesn't leak
        unset_hf_deepspeed_config()

192
193
    def test_init_zero3_fp16(self):
        # test that zero.Init() works correctly under zero3/fp16
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
        ds_config = {
            "train_batch_size": 1,
            "zero_optimization": {
                "stage": 3,
            },
        }

        dschf = HfDeepSpeedConfig(ds_config)

        self.assertTrue(dschf.is_zero3())
        self.assertTrue(is_deepspeed_zero3_enabled())

        with LoggingLevel(logging.INFO):
            with mockenv_context(**self.dist_env_1_gpu):
                logger = logging.get_logger("transformers.modeling_utils")
                with CaptureLogger(logger) as cl:
                    AutoModel.from_pretrained(T5_TINY)
        self.assertIn("Detected DeepSpeed ZeRO-3", cl.out)

        # now remove zero optimization
        del ds_config["zero_optimization"]
        dschf = HfDeepSpeedConfig(ds_config)

        self.assertFalse(dschf.is_zero3())
        self.assertFalse(is_deepspeed_zero3_enabled())

        with LoggingLevel(logging.INFO):
            with mockenv_context(**self.dist_env_1_gpu):
                logger = logging.get_logger("transformers.modeling_utils")
                with CaptureLogger(logger) as cl:
                    AutoModel.from_pretrained(T5_TINY)
        self.assertNotIn("Detected DeepSpeed ZeRO-3", cl.out)


228
class TrainerIntegrationDeepSpeedWithCustomConfig(TestCasePlus):
229
230
    def setUp(self):
        super().setUp()
231
232
233
234
235

        args = TrainingArguments(".")
        self.n_epochs = args.num_train_epochs
        self.batch_size = args.train_batch_size

Stas Bekman's avatar
Stas Bekman committed
236
        master_port = get_master_port(real_launcher=False)
237
238
239
240
241
242
243
        self.dist_env_1_gpu = {
            "MASTER_ADDR": "localhost",
            "MASTER_PORT": master_port,
            "RANK": "0",
            "LOCAL_RANK": "0",
            "WORLD_SIZE": "1",
        }
244

245
246
247
248
        self.ds_config_file = {
            "zero2": f"{self.test_file_dir_str}/ds_config_zero2.json",
            "zero3": f"{self.test_file_dir_str}/ds_config_zero3.json",
        }
249
250
251

        # use self.get_config_dict(stage) to use these to ensure the original is not modified
        with io.open(self.ds_config_file[ZERO2], "r", encoding="utf-8") as f:
252
            config_zero2 = json.load(f)
253
        with io.open(self.ds_config_file[ZERO3], "r", encoding="utf-8") as f:
254
            config_zero3 = json.load(f)
255
            # The following setting slows things down, so don't enable it by default unless needed by a test.
256
            # It's in the file as a demo for users since we want everything to work out of the box even if slower.
257
258
            config_zero3["zero_optimization"]["stage3_gather_16bit_weights_on_model_save"] = False

259
260
261
262
        self.ds_config_dict = {
            "zero2": config_zero2,
            "zero3": config_zero3,
        }
263

264
265
266
267
268
269
    def tearDown(self):
        super().tearDown()

        # reset the ds config global so that tests state doesn't leak
        unset_hf_deepspeed_config()

270
271
272
    def get_config_dict(self, stage):
        # As some tests modify the dict, always make a copy
        return deepcopy(self.ds_config_dict[stage])
273

274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296

@require_deepspeed
@require_torch_gpu
class TrainerIntegrationDeepSpeed(TrainerIntegrationDeepSpeedWithCustomConfig, TrainerIntegrationCommon):
    """

    This class is for testing directly via get_regression_trainer

    It mixes in `TrainerIntegrationCommon` which already has a lot of helper validation methods
    which we can re-use here.

    Important: this class' setup can only work with a single gpu because it runs within the current
    pytest worker. For multi-gpu tests use TestDeepSpeedWithLauncher.

    Note: if any of the tests of this class get run there will be at least one gpu occupied by them
    until this pytest worker exits. This is because the gpu memory allocated by the cuda-kernels
    won't be released until this pytest worker exits.

    This may appear as some run-away tests if you watch `nvidia-smi` while other tests that fork new
    processes are run. So there will be one or two "stale" processes reported in `nvidia-smi`. This
    is not a bug.
    """

297
    # --- These tests are enough to run on one of zero stages --- #
298

299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
    def test_hf_ds_config_mismatch(self):
        ds_config = self.get_config_dict(ZERO2)

        # Purposefully configure these values to mismatch TrainingArguments values.
        # This currently doesn't cover all keys (but it could)
        per_device_train_batch_size = 2
        ds_config["train_micro_batch_size_per_gpu"] = per_device_train_batch_size + 2

        ds_config["train_batch_size"] = 1000

        gradient_accumulation_steps = 2
        ds_config["gradient_accumulation_steps"] = gradient_accumulation_steps + 2

        max_grad_norm = 1.0
        ds_config["gradient_clipping"] = max_grad_norm + 0.1

        adam_beta1, adam_beta2 = 0.9, 0.99
        ds_config["optimizer"]["params"]["betas"] = [adam_beta1 - 0.1, adam_beta2 - 0.1]

        fp16 = True
        ds_config["fp16"]["enabled"] = not fp16

        keys = [
            "per_device_train_batch_size",
            "train_batch_size",
            "gradient_accumulation_steps",
            "max_grad_norm",
            "betas",
            "fp16",
        ]

        with mockenv_context(**self.dist_env_1_gpu):
            trainer = get_regression_trainer(
                local_rank=0,
                fp16=fp16,
                deepspeed=ds_config,
                per_device_train_batch_size=per_device_train_batch_size,
                gradient_accumulation_steps=gradient_accumulation_steps,
                max_grad_norm=max_grad_norm,
                adam_beta1=adam_beta1,
                adam_beta2=adam_beta2,
            )
            with self.assertRaises(Exception) as context:
                trainer.train()

        for key in keys:
            self.assertTrue(
                key in str(context.exception),
                f"{key} is not in the exception message:\n{context.exception}",
            )

350
351
352
353
354
355
356
357
358
    # Test various combos
    # 1. DS scheduler + DS optimizer: this is already tested by most other tests
    # 2. HF scheduler + HF optimizer:
    # 3. DS scheduler + HF optimizer:
    # 4. HF scheduler + DS optimizer:

    def test_hf_scheduler_hf_optimizer(self):
        a = 0
        with mockenv_context(**self.dist_env_1_gpu):
359
360
361
            ds_config_zero2_dict = self.get_config_dict(ZERO2)
            del ds_config_zero2_dict["optimizer"]  # force default HF Trainer optimizer
            del ds_config_zero2_dict["scheduler"]  # force default HF Trainer scheduler
362
            ds_config_zero2_dict["zero_optimization"]["offload_optimizer"]["device"] = "none"
363
            ds_config_zero2_dict["fp16"]["initial_scale_power"] = 1  # force optimizer on the first step
364
            trainer = get_regression_trainer(a=a, local_rank=0, fp16=True, deepspeed=ds_config_zero2_dict)
365
366
367
368
369
370
371
            trainer.train()
        new_a = trainer.model.a.item()
        self.assertNotEqual(new_a, a)

    def test_ds_scheduler_hf_optimizer(self):
        a = 0
        with mockenv_context(**self.dist_env_1_gpu):
372
373
            ds_config_zero2_dict = self.get_config_dict(ZERO2)
            del ds_config_zero2_dict["optimizer"]  # force default HF Trainer optimizer
374
            ds_config_zero2_dict["zero_optimization"]["offload_optimizer"]["device"] = "none"
375
            ds_config_zero2_dict["fp16"]["initial_scale_power"] = 1  # force optimizer on the first step
376
            trainer = get_regression_trainer(a=a, local_rank=0, fp16=True, deepspeed=ds_config_zero2_dict)
377
378
379
380
381
382
            trainer.train()
        new_a = trainer.model.a.item()
        self.assertNotEqual(new_a, a)

    def test_hf_scheduler_ds_optimizer(self):
        with mockenv_context(**self.dist_env_1_gpu):
383
384
            ds_config_zero2_dict = self.get_config_dict(ZERO2)
            del ds_config_zero2_dict["scheduler"]  # force default HF Trainer scheduler
385
            ds_config_zero2_dict["zero_optimization"]["offload_optimizer"]["device"] = "none"
386
            ds_config_zero2_dict["fp16"]["initial_scale_power"] = 1  # force optimizer on the first step
387
            trainer = get_regression_trainer(local_rank=0, fp16=True, deepspeed=ds_config_zero2_dict)
388
389
390
391
392
393
394
            with self.assertRaises(Exception) as context:
                trainer.train()
        self.assertIn(
            "Found `optimizer` configured in the DeepSpeed config, but no `scheduler`. "
            "Please configure a scheduler in the DeepSpeed config.",
            str(context.exception),
        )
395

396
    @require_deepspeed_aio
397
398
399
400
401
    def test_stage3_nvme_offload(self):
        with mockenv_context(**self.dist_env_1_gpu):
            # this actually doesn't have to be on NVMe, any storage will do since this test only
            # runs a simple check that we can use some directory as if it were NVMe
            nvme_path = self.get_auto_remove_tmp_dir()
402
            nvme_config = {"device": "nvme", "nvme_path": nvme_path}
403
404
405
            ds_config_zero3_dict = self.get_config_dict(ZERO3)
            ds_config_zero3_dict["zero_optimization"]["offload_optimizer"] = nvme_config
            ds_config_zero3_dict["zero_optimization"]["offload_param"] = nvme_config
406
            trainer = get_regression_trainer(local_rank=0, fp16=True, deepspeed=ds_config_zero3_dict)
407
            with CaptureLogger(deepspeed_logger) as cl:
408
                trainer.train()
409
            self.assertIn("DeepSpeed info", cl.out, "expected DeepSpeed logger output but got none")
410

411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
    @require_optuna
    def test_hyperparameter_search(self):
        with mockenv_context(**self.dist_env_1_gpu):
            ds_config_zero3_dict = self.get_config_dict(ZERO3)

            # hyperparameter_search requires model_init() to recreate the model for each trial
            def model_init():
                config = RegressionModelConfig(a=0, b=0, double_output=False)
                model = RegressionPreTrainedModel(config)
                return model

            trainer = get_regression_trainer(
                local_rank=0,
                fp16=True,
                model_init=model_init,
                deepspeed=ds_config_zero3_dict,
            )

            n_trials = 3
            with CaptureLogger(deepspeed_logger) as cl:
                with CaptureStd() as cs:
                    trainer.hyperparameter_search(direction="maximize", n_trials=n_trials)
            self.assertIn("DeepSpeed info", cl.out, "expected DeepSpeed logger output but got none")
            self.assertIn(f"Trial {n_trials-1} finished with value", cs.err, "expected hyperparameter_search output")
            self.assertIn("Best is trial", cs.err, "expected hyperparameter_search output")

437
438
    # --- These tests need to run on both zero stages --- #

439
440
    @parameterized.expand(params, name_func=parameterized_custom_name_func)
    def test_hf_optimizer_with_offload(self, stage, dtype):
441
        # non-DS optimizers can be used with ZERO-offload (as long as they have both CPU and GPU implementation (except LAMB))
442
443
444
        ds_config_dict = self.get_config_dict(stage)
        del ds_config_dict["optimizer"]  # force default HF Trainer optimizer
        # force cpu offload
445
        ds_config_dict["zero_optimization"]["offload_optimizer"]["device"] = "cpu"
446
        ds_config_dict["zero_force_ds_cpu_optimizer"] = False  # offload is not efficient w/o CPUAdam
447
        with mockenv_context(**self.dist_env_1_gpu):
448
            kwargs = {"local_rank": 0, "deepspeed": ds_config_dict}
449
450
            kwargs[dtype] = True
            trainer = get_regression_trainer(**kwargs)
451
            with CaptureLogger(deepspeed_logger) as cl:
452
                trainer.train()
453
            self.assertIn("DeepSpeed info", cl.out, "expected DeepSpeed logger output but got none")
454

455
456
    @parameterized.expand(params, name_func=parameterized_custom_name_func)
    def test_fake_notebook_no_launcher(self, stage, dtype):
457
458
459
460
461
        # this setup emulates a notebook where a launcher needs to be emulated by hand

        # note that unittest resets sys.stdout each test, so `CaptureStd` will work here to capture
        # DeepSpeed log if this test happens to run first in this pytest worker. But it will fail if
        # it's run not as a first test as `sys.stdout` will no longer be the same. So we either have
462
463
        # to reset `deepspeed_logger.handlers[0].setStream(sys.stdout)` or directly capture from the deepspeed_logger.
        with mockenv_context(**self.dist_env_1_gpu):
464
            kwargs = {"local_rank": 0, "deepspeed": self.get_config_dict(stage)}
465
466
467
            kwargs[dtype] = True
            trainer = get_regression_trainer(**kwargs)

468
            with CaptureLogger(deepspeed_logger) as cl:
469
                trainer.train()
470
            self.assertIn("DeepSpeed info", cl.out, "expected DeepSpeed logger output but got none")
471

472
473
    @parameterized.expand(params, name_func=parameterized_custom_name_func)
    def test_early_get_last_lr(self, stage, dtype):
474
475
476
477
478
479
480
481
        # with deepspeed's fp16 and dynamic loss scale enabled the optimizer/scheduler steps may
        # not run for the first few dozen steps while loss scale is too large, and thus during
        # that time `get_last_lr` will fail if called during that warm up stage,
        #
        # setting `logging_steps=1` forces an early `trainer._maybe_log_save_evaluate()` which calls
        # `self.lr_scheduler.get_last_lr()` and originally it'd fail on the very first step.
        with mockenv_context(**self.dist_env_1_gpu):
            a = b = 0.0
482
483
484
485
486
487
488
489
490
            kwargs = {
                "a": a,
                "b": b,
                "local_rank": 0,
                "train_len": 8,
                "deepspeed": self.get_config_dict(stage),
                "per_device_train_batch_size": 8,
                "logging_steps": 1,
            }
491
492
493
            kwargs[dtype] = True
            trainer = get_regression_trainer(**kwargs)

494
            trainer.train()
495
496
            post_train_a = trainer.model.a.item()

497
498
            # XXX: for some reason the following check fails with zero3/fp16 and any/bf16 - not a
            # broken but a different qualitative outcome - as if optimizer did run
499
500
501
502
            # oddly getting 1.0 for both a and b from 0.0 - there is a bug somewhere
            # print(trainer.model.a.item())
            # print(trainer.model.b.item())
            # need to investigate at some point
503
            if (stage == ZERO3 and dtype == FP16) or (dtype == BF16):
504
                return
505
506
507

            # it's enough that train didn't fail for this test, but we must check that
            # optimizer/scheduler didn't run (since if it did this test isn't testing the right thing)
508
            self.assertEqual(post_train_a, a)
509

510
511
    @parameterized.expand(params, name_func=parameterized_custom_name_func)
    def test_gradient_accumulation(self, stage, dtype):
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
        # this test measures that we get identical weights and similar loss with:
        # 1. per_device_train_batch_size=8, gradient_accumulation_steps=1
        # 2. per_device_train_batch_size=4, gradient_accumulation_steps=2
        # since the 2nd should produce the effective batch of 1st, with the same results
        #
        # I can get an identical loss for a small train_len=32, plus the power of the initial
        # dynamic loss scale value set to:
        #   "fp16.initial_scale_power": 1
        # plus having the same WarmupLR's warmup_min_lr == warmup_max_lr in the config file
        # but for some reason going to train_len=64 the weights, weights start to mismatch with this setup.
        # the culprit seems to be `initial_scale_power` - putting it back to its default 32 keeps the weights identical

        train_len = 64
        a = b = 0.0

527
528
529
530
531
532
533
        kwargs = {
            "a": a,
            "b": b,
            "local_rank": 0,
            "train_len": train_len,
            "deepspeed": self.get_config_dict(stage),
        }
534
        kwargs[dtype] = True
535

536
537
        with mockenv_context(**self.dist_env_1_gpu):
            no_grad_accum_trainer = get_regression_trainer(
538
539
                **kwargs,
                per_device_train_batch_size=16,
540
541
542
543
544
545
546
547
548
549
550
                gradient_accumulation_steps=1,
            )
            no_grad_accum_result = no_grad_accum_trainer.train()
            no_grad_accum_loss = no_grad_accum_result.training_loss
            no_grad_accum_a = no_grad_accum_trainer.model.a.item()
            no_grad_accum_b = no_grad_accum_trainer.model.b.item()
            # make sure the optimizer kicked in - if it hasn't changed from the original value of a then make train_len bigger
            self.assertNotEqual(no_grad_accum_a, a)

        with mockenv_context(**self.dist_env_1_gpu):
            yes_grad_accum_trainer = get_regression_trainer(
551
                **kwargs,
552
                per_device_train_batch_size=4,
553
                gradient_accumulation_steps=4,
554
555
556
557
558
559
560
            )
            yes_grad_accum_result = yes_grad_accum_trainer.train()
            yes_grad_accum_loss = yes_grad_accum_result.training_loss
            yes_grad_accum_a = yes_grad_accum_trainer.model.a.item()
            yes_grad_accum_b = yes_grad_accum_trainer.model.b.item()
            self.assertNotEqual(yes_grad_accum_a, a)

561
562
563
564
        # training with half the batch size but accumulation steps as 2 should give the same
        # weights, but sometimes get a slight difference still of 1e-6
        self.assertAlmostEqual(no_grad_accum_a, yes_grad_accum_a, places=5)
        self.assertAlmostEqual(no_grad_accum_b, yes_grad_accum_b, places=5)
565
566

        # see the note above how to get identical loss on a small bs
567
        self.assertAlmostEqual(no_grad_accum_loss, yes_grad_accum_loss, places=2)
568

569
    def check_saved_checkpoints_deepspeed(self, output_dir, freq, total, stage, dtype):
570
571
572
        # adapted from TrainerIntegrationCommon.check_saved_checkpoints

        file_list = [WEIGHTS_NAME, "training_args.bin", "trainer_state.json", "config.json"]
573
574
575
576
577
578
579
580

        if stage == ZERO2:
            ds_file_list = ["mp_rank_00_model_states.pt"]
        elif stage == ZERO3:
            ds_file_list = ["zero_pp_rank_0_mp_rank_00_model_states.pt"]
        else:
            raise ValueError(f"unknown stage {stage}")

581
582
        if dtype == "bf16":
            ds_file_list.append("bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt")
583
584
585

        for step in range(freq, total, freq):
            checkpoint = os.path.join(output_dir, f"checkpoint-{step}")
586
            self.assertTrue(os.path.isdir(checkpoint), f"[{stage}] {checkpoint} dir is not found")
587
588
589

            # common files
            for filename in file_list:
590
591
                path = os.path.join(checkpoint, filename)
                self.assertTrue(os.path.isfile(path), f"[{stage}] {path} is not found")
592
593
594
595
596
597

            # ds files
            ds_path = os.path.join(checkpoint, f"global_step{step}")
            for filename in ds_file_list:
                # filename = os.path.join(path, filename)
                # print(filename)
598
599
                path = os.path.join(ds_path, filename)
                self.assertTrue(os.path.isfile(path), f"[{stage}] {path} is not found")
600

601
602
    @parameterized.expand(params, name_func=parameterized_custom_name_func)
    def test_save_checkpoints(self, stage, dtype):
603
604
        # adapted from  TrainerIntegrationTest.test_save_checkpoints

605
        freq = 5
606
        output_dir = self.get_auto_remove_tmp_dir()
607
        ds_config_dict = self.get_config_dict(stage)
608
609
610
        if dtype == FP16:
            ds_config_dict["fp16"]["initial_scale_power"] = 1  # force optimizer on the first step
        # XXX:
611
        if stage == ZERO3:
612
            ds_config_dict["zero_optimization"]["stage3_gather_16bit_weights_on_model_save"] = True
613
614
615

        # save checkpoints
        with mockenv_context(**self.dist_env_1_gpu):
616
617
618
619
620
            kwargs = {
                "output_dir": output_dir,
                "save_steps": freq,
                "deepspeed": ds_config_dict,
            }
621
622
            kwargs[dtype] = True
            trainer = get_regression_trainer(**kwargs)
623
624
625
            trainer.train()

        total = int(self.n_epochs * 64 / self.batch_size)
626
        self.check_saved_checkpoints_deepspeed(output_dir, freq, total, stage, dtype)
627

628
629
    @parameterized.expand(params, name_func=parameterized_custom_name_func)
    def test_can_resume_training_errors(self, stage, dtype):
630
631
632
        with mockenv_context(**self.dist_env_1_gpu):
            ds_config_dict = self.get_config_dict(stage)
            output_dir = self.get_auto_remove_tmp_dir()
633
            kwargs = {"output_dir": output_dir, "deepspeed": ds_config_dict}
634
635
            kwargs[dtype] = True
            trainer = get_regression_trainer(**kwargs)
636
637
638
639
640
641
642
643

            # 1. fail to find any checkpoint - due a fresh output_dir
            with self.assertRaises(Exception) as context:
                trainer.train(resume_from_checkpoint=True)
            self.assertTrue(
                "No valid checkpoint found in output directory" in str(context.exception),
                f"got exception: {context.exception}",
            )
644

645
646
647
648
649
650
651
652
            # 2. fail to find a bogus checkpoint
            with self.assertRaises(Exception) as context:
                checkpoint = os.path.join(output_dir, "checkpoint-5")
                trainer.train(resume_from_checkpoint=f"{checkpoint}-bogus")
            self.assertTrue(
                "Can't find a valid checkpoint at" in str(context.exception), f"got exception: {context.exception}"
            )

653
654
    @parameterized.expand(params_with_optims_and_schedulers, name_func=parameterized_custom_name_func)
    def test_can_resume_training_normal(self, stage, dtype, optim, scheduler):
655
656
        # adapted from TrainerIntegrationTest.test_can_resume_training
        # test normal resume for each stage separately, error-handling is tested in a different test
657
658
659
660
661
662

        # ToDo: Currently, hf_optim + hf_scheduler resumes with the correct states and
        # also has same losses for few steps but then slowly diverges. Need to figure it out.
        if optim == HF_OPTIM and scheduler == HF_SCHEDULER:
            return

663
        output_dir = self.get_auto_remove_tmp_dir("./xxx", after=False)
664
        ds_config_dict = self.get_config_dict(stage)
665
666
667
        if dtype == FP16:
            ds_config_dict["fp16"]["initial_scale_power"] = 1  # force optimizer on the first step
        # XXX:
668
        if stage == ZERO3:
669
            ds_config_dict["zero_optimization"]["stage3_gather_16bit_weights_on_model_save"] = True
670

671
672
673
674
675
676
        if optim == HF_OPTIM:
            del ds_config_dict["optimizer"]

        if scheduler == HF_SCHEDULER:
            del ds_config_dict["scheduler"]

677
678
679
680
681
682
683
        kwargs = {
            "output_dir": output_dir,
            "train_len": 128,
            "save_steps": 5,
            "learning_rate": 0.1,
            "deepspeed": ds_config_dict,
        }
684
        kwargs[dtype] = True
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716

        with mockenv_context(**self.dist_env_1_gpu):
            trainer = get_regression_trainer(**kwargs)
            trainer.train()
            (a, b) = trainer.model.a.item(), trainer.model.b.item()
            state = dataclasses.asdict(trainer.state)

            checkpoint = os.path.join(output_dir, "checkpoint-5")

            # Reinitialize trainer
            trainer = get_regression_trainer(**kwargs)

            trainer.train(resume_from_checkpoint=checkpoint)
            (a1, b1) = trainer.model.a.item(), trainer.model.b.item()
            state1 = dataclasses.asdict(trainer.state)
            self.assertEqual(a, a1)
            self.assertEqual(b, b1)
            self.check_trainer_state_are_the_same(state, state1)

            # Now check with a later checkpoint that it also works when we span over one epoch
            checkpoint = os.path.join(output_dir, "checkpoint-15")

            # Reinitialize trainer and load model
            trainer = get_regression_trainer(**kwargs)

            trainer.train(resume_from_checkpoint=checkpoint)
            (a1, b1) = trainer.model.a.item(), trainer.model.b.item()
            state1 = dataclasses.asdict(trainer.state)
            self.assertEqual(a, a1)
            self.assertEqual(b, b1)
            self.check_trainer_state_are_the_same(state, state1)

717
718
719
720
721
            # Finally, should be able to resume with the same trainer/same deepspeed engine instance
            # XXX: but currently this not possible due DS bug: https://github.com/microsoft/DeepSpeed/issues/1612
            # trainer.train(resume_from_checkpoint=checkpoint)
            # a workaround needs to be used that re-creates the deepspeed engine

722
723
    @parameterized.expand(params, name_func=parameterized_custom_name_func)
    def test_load_state_dict_from_zero_checkpoint(self, stage, dtype):
724
725
726
727
728
729
        # test that we can load fp32 weights directly from the zero checkpoint into the current model

        output_dir = self.get_auto_remove_tmp_dir()  # "./xxx", after=False, before=False)

        ds_config_dict = self.get_config_dict(stage)

730
731
732
733
734
735
736
737
738
739
        kwargs = {
            "output_dir": output_dir,
            "train_len": 4,
            "per_device_train_batch_size": 4,
            "num_train_epochs": 1,
            "save_strategy": "steps",
            "save_steps": 1,
            "learning_rate": 0.1,
            "deepspeed": ds_config_dict,
        }
740
        kwargs[dtype] = True
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756

        with mockenv_context(**self.dist_env_1_gpu):
            trainer = get_regression_trainer(**kwargs)
            trainer.train()
            (a, b) = trainer.model.a.item(), trainer.model.b.item()
            state = dataclasses.asdict(trainer.state)

            checkpoint_dir = get_last_checkpoint(output_dir)
            model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)

            (a1, b1) = model.a.item(), model.b.item()
            state1 = dataclasses.asdict(trainer.state)
            self.assertEqual(a, a1)
            self.assertEqual(b, b1)
            self.check_trainer_state_are_the_same(state, state1)

757
758
759
760
    def test_config_object(self):
        # test that we can switch from zero2 to zero3 in the same process for example
        # test is_zero, etc.
        output_dir = self.get_auto_remove_tmp_dir()
761
        kwargs = {"output_dir": output_dir, "train_len": 8, "fp16": True}
762

763
764
        ds_config_zero3_dict = self.get_config_dict(ZERO3)
        ds_config_zero2_dict = self.get_config_dict(ZERO2)
765

766
        with mockenv_context(**self.dist_env_1_gpu):
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
            trainer = get_regression_trainer(deepspeed=ds_config_zero3_dict, **kwargs)
            self.assertTrue(is_deepspeed_zero3_enabled())

            # test we can repeat that and with train this time
            trainer = get_regression_trainer(deepspeed=ds_config_zero3_dict, **kwargs)
            trainer.train()
            self.assertTrue(is_deepspeed_zero3_enabled())

            # test zero3 is disabled
            trainer = get_regression_trainer(deepspeed=ds_config_zero2_dict, **kwargs)
            self.assertFalse(is_deepspeed_zero3_enabled())

            # check config obj
            config = deepspeed_config()
            self.assertTrue(bool(config), "Deepspeed config should be accessible")

783
784
            # with accelerate integration below line is additionally required for this test to pass
            trainer.accelerator.state._reset_state()
785
786
787
788
789
790
            del trainer
            # now weakref should gc the global and we shouldn't get anything here
            config = deepspeed_config()
            self.assertFalse(is_deepspeed_zero3_enabled())
            self.assertFalse(bool(config), "Deepspeed config should not be accessible")

791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
    @parameterized.expand(params, name_func=parameterized_custom_name_func)
    def test_load_best_model(self, stage, dtype):
        # Test that forced deepspeed reinit doesn't break the model. the forced re-init after
        # loading the best model in Trainer is there to workaround this bug in Deepspeed
        # https://github.com/microsoft/DeepSpeed/issues/1612
        #
        # The test is derived from a repro script submitted in this Issue:
        # https://github.com/huggingface/transformers/issues/17114
        #
        # One additional feature of this test is that we use a non-AdamW optimizer to test that
        # deepspeed doesn't fallback to AdamW, which would prevent the optimizer states from loading
        # correctly

        from transformers import T5ForConditionalGeneration, T5Tokenizer, Trainer  # noqa

        output_dir = self.get_auto_remove_tmp_dir()  # "./xxx", after=False, before=False)

        ds_config_dict = self.get_config_dict(stage)
        del ds_config_dict["optimizer"]  # will use HF Trainer optimizer
        del ds_config_dict["scheduler"]  # will use HF Trainer scheduler
811
        ds_config_dict["zero_force_ds_cpu_optimizer"] = False  # offload is not efficient w/o CPUAdam
812
813
814
        # must use this setting to get the reload path exercised
        ds_config_dict["zero_optimization"]["stage3_gather_16bit_weights_on_model_save"] = True

815
        with mockenv_context(**self.dist_env_1_gpu):
816
            args_dict = {
817
818
                "per_device_train_batch_size": 1,
                "per_device_eval_batch_size": 1,
819
820
821
822
823
824
825
826
827
828
829
830
831
                "gradient_accumulation_steps": 1,
                "learning_rate": 1e-4,
                "num_train_epochs": 1,
                "do_train": True,
                "do_eval": True,
                "optim": "adafactor",
                "evaluation_strategy": "steps",
                "eval_steps": 1,
                "save_strategy": "steps",
                "save_steps": 1,
                "load_best_model_at_end": True,
                "max_steps": 1,
                "deepspeed": ds_config_dict,
832
                "report_to": "none",
833
834
835
            }

            training_args = TrainingArguments(output_dir, **args_dict)
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
            tokenizer = T5Tokenizer.from_pretrained(T5_TINY)
            model = T5ForConditionalGeneration.from_pretrained(T5_TINY)

            def _add_eos_to_examples(example):
                example["input_text"] = f"question: {example['question']}  context: {example['context']}"
                example["target_text"] = example["answers"]["text"][0] if len(example["answers"]["text"]) > 0 else ""
                return example

            def _convert_to_features(example_batch):
                input_encodings = tokenizer.batch_encode_plus(
                    example_batch["input_text"], pad_to_max_length=True, max_length=512, truncation=True
                )
                target_encodings = tokenizer.batch_encode_plus(
                    example_batch["target_text"], pad_to_max_length=True, max_length=16, truncation=True
                )

                encodings = {
                    "input_ids": input_encodings["input_ids"],
                    "attention_mask": input_encodings["attention_mask"],
                    "labels": target_encodings["input_ids"],
                }

                return encodings

            def get_dataset():
                data_file = str(self.tests_dir / "fixtures/tests_samples/SQUAD/sample.json")
862
                data_files = {"train": data_file, "validation": data_file}
863
864
865
866
867
868
869
                raw_datasets = datasets.load_dataset("json", data_files=data_files, field="data")
                train_dataset = raw_datasets["train"].map(_add_eos_to_examples).map(_convert_to_features, batched=True)
                valid_dataset = deepcopy(train_dataset)
                return train_dataset, valid_dataset

            train_dataset, eval_dataset = get_dataset()

870
871
872
873
874
875
876
877
878
879
            trainer = Trainer(
                model=model,
                tokenizer=tokenizer,
                args=training_args,
                train_dataset=train_dataset,
                eval_dataset=eval_dataset,
            )
            trainer.train()  # crash 1 was here
            trainer.evaluate()  # crash 2 was here

880
881
882
883

@slow
@require_deepspeed
@require_torch_gpu
884
class TestDeepSpeedWithLauncher(TestCasePlus):
Patrick von Platen's avatar
Patrick von Platen committed
885
    """This class is for testing via an external script - can do multiple gpus"""
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901

    # Tests to devise #
    #
    # 1. predict_with_generate on multigpu - need to figure out how to give input sequences so that
    # the 2 gpus will generate prediction sequences that aren't of the same length - this is because
    # we had to code a special feature to sync the gpus when the predicted sequences aren't of the
    # same length. In general this will tested as a side-effect through a variety of other tests -
    # it'll simply hang trying to synchronize with other gpus if this problem is encountered. So as
    # long as we have a few full tests running on zero3 + predict_with_generate this should be
    # mostly covered.
    #
    # but there are 5 variations on beam search in `generate`- with identical code branched with `if
    # synced_gpus`
    #
    # 2. most tests should probably be run on both: zero2 and zero3 configs
    #
902

903
    @parameterized.expand(params, name_func=parameterized_custom_name_func)
Yih-Dar's avatar
Yih-Dar committed
904
    @require_torch_multi_gpu
905
906
    def test_basic_distributed(self, stage, dtype):
        self.run_and_check(stage=stage, dtype=dtype, distributed=True)
907

908
909
    def test_do_eval_no_train(self):
        # testing only zero3 since zero2 makes no sense with inference
910
        self.run_and_check(
911
            stage=ZERO3,
912
            dtype=FP16,
913
914
            eval_steps=1,
            distributed=False,
915
916
            do_train=False,
            do_eval=True,
917
        )
918

919
920
    @parameterized.expand(params, name_func=parameterized_custom_name_func)
    def test_fp32_non_distributed(self, stage, dtype):
921
922
923
924
        # real model needs too much GPU memory under stage2+fp32, so using tiny random model here -
        # therefore no quality checks, just basic completion checks are done
        self.run_and_check(
            stage=stage,
925
            dtype=dtype,
926
927
928
929
930
            model_name=T5_TINY,
            distributed=False,
            do_train=True,
            do_eval=True,
            quality_checks=False,
931
            fp32=True,
932
933
        )

934
    @parameterized.expand(params, name_func=parameterized_custom_name_func)
Yih-Dar's avatar
Yih-Dar committed
935
    @require_torch_multi_gpu
936
    def test_fp32_distributed(self, stage, dtype):
937
938
939
940
        # real model needs too much GPU memory under stage2+fp32, so using tiny random model here -
        # therefore no quality checks, just basic completion checks are done
        self.run_and_check(
            stage=stage,
941
            dtype=dtype,
942
943
944
945
946
            model_name=T5_TINY,
            distributed=True,
            do_train=True,
            do_eval=True,
            quality_checks=False,
947
            fp32=True,
948
949
        )

950
951
    @parameterized.expand(params, name_func=parameterized_custom_name_func)
    def test_resume_train_not_from_ds_checkpoint(self, stage, dtype):
952
953
954
955
956
        # do normal training and then resume not from the deepspeed checkpoint but explicitly from
        # the saved model dir

        do_train = True
        do_eval = False
957
958
959
960
961
962
963
964
        kwargs = {
            "stage": stage,
            "dtype": dtype,
            "eval_steps": 1,
            "distributed": True,
            "do_train": do_train,
            "do_eval": do_eval,
        }
965
966
967
968
969
970
971
972
973
974

        # 1. normal training
        output_dir = self.run_and_check(**kwargs)

        # 2. now resume explicitly from the saved weights, by passing --model_name_or_path output_dir
        # - i.e. the same path the model was saved to in step 1
        output_dir = self.run_trainer(**kwargs, model_name=output_dir)

        self.do_checks(output_dir, do_train=do_train, do_eval=do_eval)

975
    @parameterized.expand(["bf16", "fp16", "fp32"])
Yih-Dar's avatar
Yih-Dar committed
976
    @require_torch_multi_gpu
977
    def test_inference(self, dtype):
978
        if dtype == "bf16" and not is_torch_bf16_gpu_available():
979
980
            self.skipTest("test requires bfloat16 hardware support")

981
982
        # this is just inference, so no optimizer should be loaded
        # it only works for z3 (makes no sense with z1-z2)
983
        fp32 = True if dtype == "fp32" else False
984
985
        self.run_and_check(
            stage=ZERO3,
986
            dtype=FP16,
987
988
989
990
991
            model_name=T5_TINY,
            distributed=True,
            do_train=False,
            do_eval=True,
            quality_checks=False,
992
            fp32=fp32,
993
994
        )

995
    def do_checks(self, output_dir, do_train=True, do_eval=True, quality_checks=True):
996
997
998
        if do_train:
            train_metrics = load_json(os.path.join(output_dir, "train_results.json"))
            self.assertIn("train_samples_per_second", train_metrics)
999
1000
            if quality_checks:
                self.assertGreater(train_metrics["train_samples_per_second"], 0.5)
1001
1002
1003
1004

        if do_eval:
            eval_metrics = load_json(os.path.join(output_dir, "eval_results.json"))
            self.assertIn("eval_bleu", eval_metrics)
1005
1006
            if quality_checks:
                self.assertGreater(eval_metrics["eval_bleu"], 1)
1007
1008

    # XXX: need to do better validation beyond just that the run was successful
1009
1010
1011
    def run_and_check(
        self,
        stage,
1012
        dtype,
1013
1014
1015
1016
1017
1018
        model_name: str = T5_SMALL,
        eval_steps: int = 10,
        distributed: bool = True,
        do_train: bool = True,
        do_eval: bool = True,
        quality_checks: bool = True,
1019
        fp32: bool = False,
1020
1021
        extra_args_str: str = None,
        remove_args_str: str = None,
1022
1023
    ):
        # we are doing quality testing so using a small real model
1024
        output_dir = self.run_trainer(
1025
            stage=stage,
1026
            dtype=dtype,
1027
            model_name=model_name,
1028
            eval_steps=eval_steps,
1029
            num_train_epochs=1,
1030
1031
            do_train=do_train,
            do_eval=do_eval,
1032
            distributed=distributed,
1033
            fp32=fp32,
1034
1035
1036
            extra_args_str=extra_args_str,
            remove_args_str=remove_args_str,
        )
1037

1038
        self.do_checks(output_dir, do_train=do_train, do_eval=do_eval, quality_checks=quality_checks)
1039
1040

        return output_dir
1041
1042
1043

    def run_trainer(
        self,
1044
        stage: str,
1045
        dtype: str,
1046
        model_name: str,
1047
1048
1049
1050
        eval_steps: int = 10,
        num_train_epochs: int = 1,
        do_train: bool = False,
        do_eval: bool = True,
1051
        distributed: bool = True,
1052
        fp32: bool = False,
1053
1054
1055
        extra_args_str: str = None,
        remove_args_str: str = None,
    ):
1056
        max_len = 32
Sylvain Gugger's avatar
Sylvain Gugger committed
1057
        data_dir = self.test_file_dir / "../fixtures/tests_samples/wmt_en_ro"
1058
1059
1060
        output_dir = self.get_auto_remove_tmp_dir()
        args = f"""
            --model_name_or_path {model_name}
1061
1062
            --train_file {data_dir}/train.json
            --validation_file {data_dir}/val.json
1063
1064
1065
1066
1067
1068
1069
            --output_dir {output_dir}
            --overwrite_output_dir
            --max_source_length {max_len}
            --max_target_length {max_len}
            --val_max_target_length {max_len}
            --warmup_steps 8
            --predict_with_generate
1070
1071
            --save_steps 0
            --eval_steps {eval_steps}
1072
1073
            --group_by_length
            --label_smoothing_factor 0.1
1074
1075
            --source_lang en
            --target_lang ro
1076
            --report_to none
1077
        """.split()
1078
1079
        args.extend(["--source_prefix", '"translate English to Romanian: "'])

1080
1081
        if not fp32:
            args.extend([f"--{dtype}"])
1082

1083
1084
1085
1086
1087
1088
1089
        actions = 0
        if do_train:
            actions += 1
            args.extend(
                f"""
            --do_train
            --num_train_epochs {str(num_train_epochs)}
1090
            --max_train_samples 16
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
            --per_device_train_batch_size 2
            --learning_rate 3e-3
            """.split()
            )

        if do_eval:
            actions += 1
            args.extend(
                """
            --do_eval
1101
            --max_eval_samples 16
1102
1103
1104
1105
1106
            --per_device_eval_batch_size 2
            """.split()
            )

        assert actions > 0, "need at least do_train or do_eval for the test to run"
1107
1108
1109
1110

        if extra_args_str is not None:
            args.extend(extra_args_str.split())

1111
        # currently only works for bool args
1112
1113
1114
1115
        if remove_args_str is not None:
            remove_args = remove_args_str.split()
            args = [x for x in args if x not in remove_args]

1116
        ds_args = f"--deepspeed {self.test_file_dir_str}/ds_config_{stage}.json".split()
Sylvain Gugger's avatar
Sylvain Gugger committed
1117
        script = [f"{self.examples_dir_str}/pytorch/translation/run_translation.py"]
1118
        launcher = get_launcher(distributed)
1119
1120

        cmd = launcher + script + args + ds_args
1121
        # keep for quick debug
1122
1123
1124
1125
1126
        # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
        execute_subprocess_async(cmd, env=self.get_env())

        return output_dir

1127
1128
    @parameterized.expand(params, name_func=parameterized_custom_name_func)
    def test_clm(self, stage, dtype):
1129
1130
1131
1132
1133
1134
        # this test exercises model.resize_token_embeddings() which requires param gathering outside
        # of forward - it's not used by `run_translation.py`, but it is in `run_clm.py`

        data_dir = self.tests_dir / "fixtures"
        output_dir = self.get_auto_remove_tmp_dir()
        args = f"""
1135
            --model_name_or_path {GPT2_TINY}
1136
1137
1138
1139
1140
1141
            --train_file {data_dir}/sample_text.txt
            --validation_file {data_dir}/sample_text.txt
            --output_dir {output_dir}
            --overwrite_output_dir
            --do_train
            --do_eval
1142
1143
1144
1145
            --max_train_samples 16
            --max_eval_samples 16
            --per_device_train_batch_size 2
            --per_device_eval_batch_size 2
1146
1147
            --num_train_epochs 1
            --warmup_steps 8
1148
            --block_size 64
1149
            --report_to none
1150
1151
            """.split()

1152
1153
        args.extend([f"--{dtype}"])

1154
        ds_args = f"--deepspeed {self.test_file_dir_str}/ds_config_{stage}.json".split()
Sylvain Gugger's avatar
Sylvain Gugger committed
1155
        script = [f"{self.examples_dir_str}/pytorch/language-modeling/run_clm.py"]
1156
        launcher = get_launcher(distributed=True)
1157
1158
1159
1160

        cmd = launcher + script + args + ds_args
        # keep for quick debug
        # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
1161
1162
        execute_subprocess_async(cmd, env=self.get_env())

1163
    def test_clm_from_config_zero3_fp16(self):
1164
1165
1166
1167
1168
1169
        # this test exercises AutoModel.from_config(config) - to ensure zero.Init is called

        data_dir = self.tests_dir / "fixtures"
        output_dir = self.get_auto_remove_tmp_dir()
        args = f"""
            --model_type gpt2
1170
            --tokenizer_name {GPT2_TINY}
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
            --train_file {data_dir}/sample_text.txt
            --validation_file {data_dir}/sample_text.txt
            --output_dir {output_dir}
            --overwrite_output_dir
            --do_train
            --max_train_samples 4
            --per_device_train_batch_size 2
            --num_train_epochs 1
            --warmup_steps 8
            --block_size 8
            --fp16
            --report_to none
            """.split()

        ds_args = f"--deepspeed {self.test_file_dir_str}/ds_config_zero3.json".split()
        script = [f"{self.examples_dir_str}/pytorch/language-modeling/run_clm.py"]
1187
        launcher = get_launcher(distributed=True)
1188
1189
1190
1191
1192
1193

        cmd = launcher + script + args + ds_args
        # keep for quick debug
        # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
        with CaptureStderr() as cs:
            execute_subprocess_async(cmd, env=self.get_env())
1194
        self.assertIn("Detected DeepSpeed ZeRO-3", cs.err)