Gradient Accumulation

While normally DeepSpeed gets gradient accumulation configured with:

.. code-block:: json

in this case, to enable gradient accumulation, pass the command line `--gradient_accumulation_steps` argument as normal

and it will get injected into the DeepSpeed configuration.

If you try to add it directly to the configuration file, you will receive an error from the Trainer - this is because

this setting is needed by the Trainer too, and so this approach ensures that there is a single way of setting this

value and thus avoid potential subtle errors.

"initial_scale_power": 32,

mockenv_context,

bindir = os.path.abspath(os.path.dirname(__file__))

sys.path.append(f"{bindir}/../../../tests")

from test_trainer import get_regression_trainer # noqa

class TrainerIntegrationDeepSpeed(TestCasePlus):

""" This class is for testing directly via get_regression_trainer """

def setUp(self):

super().setUp()

self.dist_env_1_gpu = dict(

MASTER_ADDR="localhost", MASTER_PORT="10999", RANK="0", LOCAL_RANK="0", WORLD_SIZE="1"

)

self.ds_config_file = f"{self.test_file_dir_str}/ds_config.json"

# this setup emulates a notebook where a launcher needs to be emulated by hand

with mockenv_context(**self.dist_env_1_gpu):

trainer = get_regression_trainer(local_rank=0, deepspeed=self.ds_config_file)

trainer.train()

def test_gradient_accumulation(self):

# 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

with mockenv_context(**self.dist_env_1_gpu):

no_grad_accum_trainer = get_regression_trainer(

a=a,

b=b,

local_rank=0,

train_len=train_len,

deepspeed=self.ds_config_file,

per_device_train_batch_size=8,

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(

a=a,

b=b,

local_rank=0,

train_len=train_len,

deepspeed=self.ds_config_file,

per_device_train_batch_size=4,

gradient_accumulation_steps=2,

)

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)

# training with half the batch size but accumulation steps as 2 should give the same weights

self.assertEqual(no_grad_accum_a, yes_grad_accum_a)

self.assertEqual(no_grad_accum_b, yes_grad_accum_b)

# see the note above how to get identical loss on a small bs

self.assertAlmostEqual(no_grad_accum_loss, yes_grad_accum_loss, places=5)

class TestDeepSpeed(TestCasePlus):

""" This class is for testing via an external script """

import contextlib

def mockenv_context(*remove, **update):

"""

env = os.environ

update = update or {}

remove = remove or []

# List of environment variables being updated or removed.

stomped = (set(update.keys()) | set(remove)) & set(env.keys())

# Environment variables and values to restore on exit.

update_after = {k: env[k] for k in stomped}

# Environment variables and values to remove on exit.

remove_after = frozenset(k for k in update if k not in env)

try:

env.update(update)

[env.pop(k, None) for k in remove]

yield

finally:

env.update(update_after)

[env.pop(k) for k in remove_after]

return self.deepspeed

# Optimizer step for deepspeed must be called on every step regardless of the value of gradient_accumulation_steps

if self.deepspeed:

self.deepspeed.step()

pass # called outside the loop

if not self.deepspeed:

self.lr_scheduler.step()

if self.args.gradient_accumulation_steps > 1 and not self.deepspeed:

# deepspeed handles loss scaling by gradient_accumulation_steps in its `backward`

# loss gets scaled under gradient_accumulation_steps in deepspeed

loss = self.deepspeed.backward(loss)