Repository: facebookresearch/fairscale

Issue:
AdaScale with gradient accumulation: Sum or average of gradients?
Hi!

I am wondering what is the intended way to implement gradient accumulation when using AdaScale. Usually, for gradient accumulation, I would use a training loop like this

    for i, (X, y) in enumerate(train_loader):
        if (i+1) % num_accum_steps == 0:
            loss = loss_fn(model(X), y) / num_accum_steps
            loss.backward()
            opt.step()
            opt.zero_grad()
        else:
            with model.no_sync():
                loss = loss_fn(model(X), y) / num_accum_steps
                loss.backward()

That is, when the optimizer step is applied, the gradient stored in `p.grad` is the **average** gradient across accumulation steps. However, this [line](https://github.com/facebookresearch/fairscale/blob/main/fairscale/optim/adascale.py#L460) in the AdaScale code makes me think that AdaScale expects `p.grad` to contain the **sum** of gradients across accumulation steps.

Could somebody clarify what's the intended way?

Thanks!


Discussion:
Sorry, I missed the question for a few days.

AdaScale treats each grad step as if it is from a separate rank. That means you can use grad-accumulation as a way to simulate multiple-gpu training with adascale.

Then, if you don't really want that, you can override the scale value with adascale's init function. I.e. if you use 8gpu and 8-step grad accumulation. Normally, adascale will think this is scale 64. But you can specify scale=8 and it should work. Let me know if it doesn't.

BTW, what are you using adascale for? (if you can share)

>  the gradient stored in p.grad is the average gradient across accumulation steps.

We usually use *sum" not average in p.grad (hence the name "accumulation").
Thanks a lot for the reply.

I don't want to use gradient accumulation to _simulate_ multi-GPU training, but in combination with it. When communication overhead becomes large, I would like to perform multiple gradient accumulation steps locally on each worker, then sync between workers and update. 

I'm confused about the sum vs average. If I understand correctly, when a model is wrapped in `DDP`, after each backward pass, the gradients in `p.grad` are **averaged** across all workers. So I would assume that the gradient accumulation should also be averaged, in order to retain the original scale of the gradient. Otherwise, if I update with a hypothetical ideal gain of `num_workers * num_accum_steps` I will overshoot since `p.grad` already contains the _sum_ of `num_accum_steps` gradients.

But that doesn't seem to square with this [line](https://github.com/facebookresearch/fairscale/blob/main/fairscale/optim/adascale.py#L460) in the code as well as your answer.
> when a model is wrapped in DDP, after each backward pass, the gradients in p.grad are averaged across all workers

Maybe I am mistaken, but I think DDP does a sum. Here is the DDP code:

https://github.com/pytorch/pytorch/blob/eff312f07be85508f049e5ddfe9ada5aa5df4fc4/torch/nn/parallel/distributed.py#L1103

then, in distributed, the default op is sum:

https://github.com/pytorch/pytorch/blob/c29b7865d02239d89d4407559a85a556039cb7c6/torch/distributed/distributed_c10d.py#L1318

Regardless it is DDP or grad-accumulation, it is doing a bigger batch, which to me makes sense to sum all grads from the entire batch.
Hi, 

Taking the code posted above 
```
for i, (X, y) in enumerate(train_loader):
    if (i+1) % num_accum_steps == 0:
        loss = loss_fn(model(X), y) / num_accum_steps    ---- averaging the loss over the #grad accum batches. 
        loss.backward()
        opt.step()
        opt.zero_grad()
    else:
        with model.no_sync():
            loss = loss_fn(model(X), y) / num_accum_steps  ---- averaging the loss over the #grad accum batches.
            loss.backward()
```


To be clear,  
- we divide the loss by the number of accumulation steps to compute the mean of the loss. (as indicated in the code above). We wouldn't do this explicitly without gradient accumulation because MSE, CE in Pytorch implement batch averaging. 
- If we are computing the mean loss, what is the [line 460](https://github.com/facebookresearch/fairscale/blob/main/fairscale/optim/adascale.py#L460) in Adascale for? 

Can you help us with specifically what we incorrectly understand here?
Thanks for persisting. I did some investigation from the git history. It seems that line 460 is a bug. It was originally introduced when grad accumulation was first added to adascale. I only verified it by training cifar, which as I now know is very tolerating of mistakes.

I went though the math again and I think the scaling on line 460 is incorrect and should be removed if gradients are accumulated.

In your case though, if you already have loss (hence gradient) pre-divided as your code is showing, then I think it is the following:

Say you have have 2 ranks and grad-accum = 2. That means you have g1 and g2 from rank 0 and g3 and g4 from rank 1.
Since you have scaled them, they became g1/2 and g2/2 and g3/2 and g4/2. For adascale, local_grad_sqr becomes:

(g1/2)^2 + (g2/2)^2 + (g3/2)^2 + (g4/2)^2 = (original local_grad_sqr) / 4

Then, the total_grad_sqr became:

((g1+g2+g3+g4)/2)^2 = (g1+g2+g3+g4)^2/4 = (original total_grad_sqr) / 4

So I think in that case, the line 460 scaling is still not needed.

I think with this bug, the gain is over-estimated when gradient accumulation is enabled. Is that what you are observing?
Thanks very much for digging into this. Highly appreciated!

I will adapt the code and run some experiments (probably towards the end of the week) and report back here.
👍thanks！
cc @mikerabbat 
@lballes @prabhuteja12, I revisited this issue and found that we actually have unit tests that ensure grad_accum and DDP are the same. See here:

![image](https://user-images.githubusercontent.com/24926999/189494381-ad66552a-3ea6-4de6-8779-3c641b6dcc17.png)


If I remove the `total_grad_sqr = total_grad_sqr / (self._num_grads_to_accum**2)` line a bunch of unit test fails. The test data is in fact shared between ddp and grad_accum tests. Therefore, this line of normalization is in fact needed unless there is some error within the tests.

You can remove the line and see the test error yourself. The errors are:

![image](https://user-images.githubusercontent.com/24926999/189494484-454a9409-2111-4cb0-85a3-14805d43965f.png)


Resolving patch:
diff --git a/fairscale/optim/adascale.py b/fairscale/optim/adascale.py
--- a/fairscale/optim/adascale.py
+++ b/fairscale/optim/adascale.py
@@ -130,7 +130,12 @@ class AdaScale(Optimizer):
             between each optimizer step. This can be changed during
             training as long as the train loop changes gradient accumulation
             accordingly.
+            The loss in each pass can be either scaled or unscaled. See `is_scaled_loss` below.
             Default to 1, which does not accumulate gradients.
+        is_scaled_loss (bool):
+            If True, assume that the loss is scaled by `num_gradients_to_accumulate`.
+            If False, the loss is not scaled.
+            Default: True.
         debias_ewma (bool):
             (experimental) Use debias exponential moving average
             for smoothing and mu and sigma variables. False will
@@ -145,6 +150,7 @@ def __init__(
         scale: Optional[float] = None,
         smoothing: float = None,
         num_gradients_to_accumulate: int = 1,
+        is_scaled_loss: bool = True,
         debias_ewma: bool = True,
     ):
         # Init hook_handles list, otherwise, a partial init'ed object may fail in ``__del__``.
@@ -160,6 +166,7 @@ def __init__(
         self._last_final_backward_call = 0
         self._num_grads_to_accum = num_gradients_to_accumulate
         self._debias_ewma = debias_ewma
+        self._is_scaled_loss = is_scaled_loss
 
         # Proxy the param_groups so that `torch.optim.lr_scheduler` can work.
         self.param_groups = self._optimizer.param_groups
@@ -453,17 +460,22 @@ def _final_callback(self) -> None:
         total_grad_sqr = np.array(
             [sum(param.grad.pow(2).sum().item() for param in group["params"]) for group in self._optimizer.param_groups]
         )
-        # Divide by (_num_grads_to_accum ** 2) to account for gradient
-        # accumulation.
-        if self._num_grads_to_accum > 1:
-            # np array doesn't support /=.
-            total_grad_sqr = total_grad_sqr / (self._num_grads_to_accum**2)
 
         # Wait for all_reduce to be done and move it to cpu & np.
         if work:
             work.wait()
         local_grad_sqr = self._local_grad_sqr.cpu().numpy()
 
+        if self._num_grads_to_accum > 1:
+            # Handle scaling for for gradient accumulation
+            if self._is_scaled_loss:
+                # If loss is scaled down, we need to scale the local gradients back by a factor of _num_grads_to_accum squared;
+                # total_grad_sqr is already scaled by _num_grads_to_accum squared.
+                local_grad_sqr *= self._num_grads_to_accum**2
+            else:
+                # If loss is not scaled, local gradients are correct, but we need to scale the total_grad_sqr down to account for gradient accumulation.
+                total_grad_sqr /= self._num_grads_to_accum**2
+
         # See appendix B.3 of the paper.
         # Modified to handle cases where scale != world_size
         #
@@ -509,7 +521,7 @@ def step(self, *args: Any, **kwargs: Any) -> Optional[float]:
         original_lr = []
         for idx, param_group in enumerate(self._optimizer.param_groups):
             original_lr.append(param_group["lr"])
-            param_group["lr"] = self.gain(pg_idx=idx) * param_group["lr"]
+            param_group["lr"] *= self.gain(pg_idx=idx)
 
         # Step it.
         res = self._optimizer.step(*args, **kwargs)
@@ -606,6 +618,18 @@ def set_num_gradients_to_accumulate(
             # not needed, so the smoothing factor is 0.
             self._smoothing = max(1 - self._world_size * self._num_grads_to_accum / 1000, 0)
 
+    def scale_grad_by_num_grads_to_accum(self) -> None:
+        """Scale the gradient down by the number of gradients to accumulate.
+
+        This should be called after the gradient accumulation is done and the unscaled loss is used.
+        """
+        assert self._local_grad_sqr is None, "Only call this after backward"
+        assert self._num_grads_to_accum > 1, "Must be accumulating gradients"
+        assert not self._is_scaled_loss, "Must use unscaled loss"
+        for group in self._optimizer.param_groups:
+            for param in group["params"]:
+                param.grad.div_(self._num_grads_to_accum)
+
     def __getattr__(self, name: str) -> Any:
         """Forward missing attributes to wrapped optimizer."""
         try:


Test patch:
diff --git a/fairscale/fair_dev/testing/golden_testing_data.py b/fairscale/fair_dev/testing/golden_testing_data.py
--- a/fairscale/fair_dev/testing/golden_testing_data.py
+++ b/fairscale/fair_dev/testing/golden_testing_data.py
@@ -8,12 +8,42 @@
 
 adascale_test_data = [
     # "input" value is a list of input tensors for micro-batch/rank 0 and micro-batch/rank 1.
-    {"input": [[1.0, 0], [0, 1.0]], "expected_gain": 4.0 / 3},
-    {"input": [[1.0, 1.0], [1.0, 1.0]], "expected_gain": 1.0000001249999846},
-    {"input": [[-1.0, 1.0], [1.0, -1.0]], "expected_gain": 2.0},
-    {"input": [[1.0, 4.0], [5.0, 0.5]], "expected_gain": 1.4688796680497926},
-    {"input": [[-0.2, 3.0], [5.0, 0.5]], "expected_gain": 1.8472893901708},
+    {
+        "input": [[1.0, 0], [0, 1.0]],
+        "expected_gain": 4.0 / 3,
+        "expected_grad": [[0.5, 0.5], [0.5, 0.5]],
+        "expected_bias_grad": [1.0, 1.0],
+    },
+    {
+        "input": [[1.0, 1.0], [1.0, 1.0]],
+        "expected_gain": 1.0000001249999846,
+        "expected_grad": [[1.0, 1.0], [1.0, 1.0]],
+        "expected_bias_grad": [1.0, 1.0],
+    },
+    {
+        "input": [[-1.0, 1.0], [1.0, -1.0]],
+        "expected_gain": 2.0,
+        "expected_grad": [[0.0, 0.0], [0.0, 0.0]],
+        "expected_bias_grad": [1.0, 1.0],
+    },
+    {
+        "input": [[1.0, 4.0], [5.0, 0.5]],
+        "expected_gain": 1.4688796680497926,
+        "expected_grad": [[3.0, 2.25], [3.0, 2.25]],
+        "expected_bias_grad": [1.0, 1.0],
+    },
+    {
+        "input": [[-0.2, 3.0], [5.0, 0.5]],
+        "expected_gain": 1.8472893901708,
+        "expected_grad": [[2.4000000953674316, 1.75], [2.4000000953674316, 1.75]],
+        "expected_bias_grad": [1.0, 1.0],
+    },
     # "inputs" to trigger multiple iteration tests, which make sure the
     # smoothing factor calculation is also covered.
-    {"inputs": [[[-0.2, 3.3], [5.2, 0.7]], [[1.0, 4.0], [3.1, 0.1]]], "expected_gain": 1.6720968158031417},
+    {
+        "inputs": [[[-0.2, 3.3], [5.2, 0.7]], [[1.0, 4.0], [3.1, 0.1]]],
+        "expected_gain": 1.6720968158031417,
+        "expected_grad": [[2.049999952316284, 2.049999952316284], [2.049999952316284, 2.049999952316284]],
+        "expected_bias_grad": [1.0, 1.0],
+    },
 ]
diff --git a/tests/optim/test_ddp_adascale.py b/tests/optim/test_ddp_adascale.py
--- a/tests/optim/test_ddp_adascale.py
+++ b/tests/optim/test_ddp_adascale.py
@@ -75,18 +75,25 @@ def _test_basic_func(rank, ddp_cls, world_size, tempfile_name, test_case):
         out.sum().backward()
         if ddp_cls is DDP:
             assert np.allclose(optim.gain(), test_case["expected_gain"]), optim.gain()
+            w, b = model.parameters()
+            assert np.allclose(w.grad.cpu(), test_case["expected_grad"]), w.grad
+            assert np.allclose(b.grad.cpu(), test_case["expected_bias_grad"]), b.grad
         optim.step()
         optim.zero_grad()
     else:
         # multiple iters
-        for in_data in test_case["inputs"]:
+        n = len(test_case["inputs"])
+        for i, in_data in enumerate(test_case["inputs"]):
             in_data = Tensor(in_data[rank]).cuda()
             out = model(in_data)
             out.sum().backward()
+            if i == n - 1 and ddp_cls is DDP:
+                assert np.allclose(optim.gain(), test_case["expected_gain"]), optim.gain()
+                w, b = model.parameters()
+                assert np.allclose(w.grad.cpu(), test_case["expected_grad"]), w.grad
+                assert np.allclose(b.grad.cpu(), test_case["expected_bias_grad"]), b.grad
             optim.step()
             optim.zero_grad()
-        if ddp_cls is DDP:
-            assert np.allclose(optim.gain(), test_case["expected_gain"]), optim.gain()
 
     dist.destroy_process_group()
 
diff --git a/tests/optim/test_single_node_adascale.py b/tests/optim/test_single_node_adascale.py
--- a/tests/optim/test_single_node_adascale.py
+++ b/tests/optim/test_single_node_adascale.py
@@ -61,7 +61,8 @@ def test_loss_accum_cpu():
 
 @pytest.mark.parametrize("cpu", [True, False])
 @pytest.mark.parametrize("test_case", adascale_test_data)
-def test_grad_accum(test_case, cpu):
+@pytest.mark.parametrize("is_scaled_loss", [True, False])
+def test_grad_accum(test_case, cpu, is_scaled_loss):
     """Test the basic functionality on CPU/GPU with gradient accumulation without DDP"""
     make_cudnn_deterministic()
     model = Linear(2, 2, bias=True)
@@ -69,7 +70,7 @@ def test_grad_accum(test_case, cpu):
         if torch.cuda.device_count() < 1:
             pytest.skip("1 GPU is required")
         model = model.cuda()
-    optim = AdaScale(SGD(model.parameters(), lr=0.1), num_gradients_to_accumulate=2)
+    optim = AdaScale(SGD(model.parameters(), lr=0.1), num_gradients_to_accumulate=2, is_scaled_loss=is_scaled_loss)
     expected_gain = test_case["expected_gain"]
     if "input" in test_case:
         data = [test_case["input"]] * 2
@@ -82,16 +83,27 @@ def test_grad_accum(test_case, cpu):
         in_data_0 = Tensor(in_data[0])
         if not cpu:
             in_data_0 = in_data_0.cuda()
-        out = model(in_data_0)
-        out.sum().backward()
+        loss = model(in_data_0).sum()
+        if is_scaled_loss:
+            loss = loss / 2
+        loss.backward()
         # grad pass 2
         in_data_1 = Tensor(in_data[1])
         if not cpu:
             in_data_1 = in_data_1.cuda()
-        out = model(in_data_1)
-        out.sum().backward()
+        loss = model(in_data_1).sum()
+        if is_scaled_loss:
+            loss = loss / 2
+        loss.backward()
+
+        if not is_scaled_loss:
+            optim.scale_grad_by_num_grads_to_accum()
+
         if exp_gain is not None:
             assert np.allclose(optim.gain(), exp_gain), optim.gain()
+            w, b = model.parameters()
+            assert np.allclose(w.grad.cpu(), test_case["expected_grad"]), w.grad
+            assert np.allclose(b.grad.cpu(), test_case["expected_bias_grad"]), b.grad
         # stepping it. Note that if we did more than 2 passes as promised by the
         # num_gradients_to_accumulate argument above, AdaScale is not be able to
         # detect that mistake for now. The result will just be wrong in that case.
@@ -110,14 +122,14 @@ def test_state_checkpointing():
     a unit test for checkpointing with DDP.
     """
     # Constants.
-    accum_steps = 3
+    num_grads_to_accum = 3
     in_dim = 5
 
     # Setup.
     def make_model_and_optim():
         model = Linear(in_dim, 2, bias=False)
         model = model.cuda()
-        optim = AdaScale(SGD(model.parameters(), lr=0.1, momentum=0.9), num_gradients_to_accumulate=accum_steps)
+        optim = AdaScale(SGD(model.parameters(), lr=0.1, momentum=0.9), num_gradients_to_accumulate=num_grads_to_accum)
         return model, optim
 
     model, optim = make_model_and_optim()
@@ -127,7 +139,7 @@ def run_a_bit(replay_data=None):
         data = []
         replay_data_idx = 0
         for _ in range(6):  # run some steps
-            for i in range(accum_steps):
+            for i in range(num_grads_to_accum):
                 if replay_data is None:
                     in_data = torch.rand(in_dim).cuda()
                     data.append(in_data)
@@ -136,7 +148,7 @@ def run_a_bit(replay_data=None):
                     replay_data_idx += 1
                 out = model(in_data)
                 out.sum().backward()
-                if i == accum_steps - 1:
+                if i == num_grads_to_accum - 1:
                     optim.step()
                     optim.zero_grad()
         return out, data
@@ -172,13 +184,14 @@ def run_a_bit(replay_data=None):
 
 def test_lr_scheduler():
     """Test AdaScale working with torch.optim.lr_scheduler."""
+    num_grads_to_accum = 3
     model = Linear(2, 2, bias=False)
-    optim = AdaScale(SGD(model.parameters(), lr=0.1), num_gradients_to_accumulate=3)
+    optim = AdaScale(SGD(model.parameters(), lr=0.1), num_gradients_to_accumulate=num_grads_to_accum)
     # We use 1, not 0.1 here since scheduler.step() is called here first.
     scheduler = LambdaLR(optim, lr_lambda=lambda epoch: 1 / 10**epoch)
     for epoch in range(3):
         for data_idx in range(10):
-            for accumulation in range(3):
+            for accumulation in range(num_grads_to_accum):
                 in_data = torch.rand(2)
                 loss = model(in_data).sum()
                 loss.backward()
@@ -194,8 +207,10 @@ def test_lr_scheduler():
 
 @skip_if_no_cuda
 @pytest.mark.parametrize("debias_ewma", [True, False])
-def test_add_param_group(debias_ewma):
-    """Test AdaScale supports add_param_group() API."""
+@pytest.mark.parametrize("is_scaled_loss", [True, False])
+def test_add_param_group(debias_ewma, is_scaled_loss):
+    """Test AdaScale supports add_param_group() API for both scaled and unscaled loss."""
+    num_grads_to_accum = 2
     model1 = Linear(2, 2, bias=True)
     with torch.no_grad():
         # make weights and bias deterministic, which is needed for
@@ -203,7 +218,12 @@ def test_add_param_group(debias_ewma):
         # parameters from other layers.
         model1.weight.copy_(Tensor([1.0, 2.0, 3.0, 4.0]).reshape(2, 2))
         model1.bias.fill_(0.1)
-    optim = AdaScale(SGD(model1.parameters(), lr=0.1), num_gradients_to_accumulate=2, debias_ewma=debias_ewma)
+    optim = AdaScale(
+        SGD(model1.parameters(), lr=0.1),
+        num_gradients_to_accumulate=2,
+        is_scaled_loss=is_scaled_loss,
+        debias_ewma=debias_ewma,
+    )
     assert len(optim._hook_handles) == 2, len(optim._hook_handles)
 
     model2 = Linear(2, 3, bias=True)
@@ -217,12 +237,19 @@ def test_add_param_group(debias_ewma):
     # make sure we can run the model.
     model = Sequential(model1, model2).cuda()
     in_data_0 = Tensor([1.0, 2.0]).cuda()
-    out = model(in_data_0)
-    out.sum().backward()
+    loss = model(in_data_0).sum()
+    if is_scaled_loss:
+        loss = loss / num_grads_to_accum
+    loss.backward()
 
     in_data_1 = Tensor([3.0, 4.0]).cuda()
-    out = model(in_data_1)
-    out.sum().backward()
+    loss = model(in_data_1).sum()
+    if is_scaled_loss:
+        loss = loss / num_grads_to_accum
+    loss.backward()
+
+    if not is_scaled_loss:
+        optim.scale_grad_by_num_grads_to_accum()
 
     # make sure the gains are right and we can step.
     # since this is the first step, debias_ewma doesn't affect the value.
@@ -244,19 +271,26 @@ def test_add_param_group(debias_ewma):
     # make sure we can run the model.
     model = Sequential(model1, model2, model3).cuda()
     in_data_0 = Tensor([1.0, 2.0]).cuda()
-    out = model(in_data_0)
-    out.sum().backward()
+    loss = model(in_data_0).sum()
+    if is_scaled_loss:
+        loss = loss / num_grads_to_accum
+    loss.backward()
 
     in_data_1 = Tensor([3.0, 4.0]).cuda()
-    out = model(in_data_1)
-    out.sum().backward()
+    loss = model(in_data_1).sum()
+    if is_scaled_loss:
+        loss = loss / num_grads_to_accum
+    loss.backward()
+
+    if not is_scaled_loss:
+        optim.scale_grad_by_num_grads_to_accum()
 
     # make sure gains are right and we can step.
     # the last PG's gain is not affected by debias_ewma since it is the first step for that PG.
-    assert np.allclose(optim.gain(), 1.1191193589460822 if debias_ewma else 1.1192783954732368), optim.gain()
-    assert np.allclose(optim.gain(0), 1.1428571880897151 if debias_ewma else 1.142857188085096), optim.gain(0)
-    assert np.allclose(optim.gain(1), 1.1167103578364508 if debias_ewma else 1.1167104954034948), optim.gain(1)
-    assert np.allclose(optim.gain(2), 1.117381091722702), optim.gain(2)
+    assert np.allclose(optim.gain(), 1.1382937715383077 if debias_ewma else 1.1391959826562015), optim.gain()
+    assert np.allclose(optim.gain(0), 1.142857206008338 if debias_ewma else 1.142857206006931), optim.gain(0)
+    assert np.allclose(optim.gain(1), 1.1116875516387468 if debias_ewma else 1.1116906378271827), optim.gain(1)
+    assert np.allclose(optim.gain(2), 1.0749164095196344), optim.gain(2)
     optim.step()
     optim.zero_grad()
 
@@ -264,31 +298,45 @@ def test_add_param_group(debias_ewma):
 @pytest.mark.parametrize(
     "test_case",
     [
-        {"new_accum": 3, "exp_gain": 1.2573902104603087},
-        {"new_accum": 6, "exp_gain": 1.0903738977361481},
-        {"new_accum": 9, "exp_gain": 1.0432658660558123},
+        {"num_grads_to_accum": 3, "exp_gain": 2.141385737279438},
+        {"num_grads_to_accum": 6, "exp_gain": 2.9927880097754036},
+        {"num_grads_to_accum": 9, "exp_gain": 3.4461759591877312},
     ],
 )
-def test_set_num_gradients_to_accumulate(test_case):
+@pytest.mark.parametrize("is_scaled_loss", [True, False])
+def test_set_num_gradients_to_accumulate(test_case, is_scaled_loss):
     """Test set_num_gradients_to_accumulate experimental feature."""
-    new_accum = test_case["new_accum"]
+    num_grads_to_accum = test_case["num_grads_to_accum"]
     exp_gain = test_case["exp_gain"]
 
     model = Linear(2, 2, bias=False)
-    optim = AdaScale(SGD(model.parameters(), lr=0.1), num_gradients_to_accumulate=2)
-    out = model(Tensor([0.0, 1.0]))
-    out.sum().backward()
-    out = model(Tensor([1.0, 0.0]))
-    out.sum().backward()
+    optim = AdaScale(SGD(model.parameters(), lr=0.1), num_gradients_to_accumulate=2, is_scaled_loss=is_scaled_loss)
+    loss = model(Tensor([0.0, 1.0])).sum()
+    if is_scaled_loss:
+        loss = loss / 2
+    loss.backward()
+    loss = model(Tensor([1.0, 0.0])).sum()
+    if is_scaled_loss:
+        loss = loss / 2
+    loss.backward()
+
+    if not is_scaled_loss:
+        optim.scale_grad_by_num_grads_to_accum()
+
     assert np.allclose(optim.gain(), 2.0)
     optim.step()
     optim.zero_grad()
 
-    optim.set_scale(float(new_accum))
-    optim.set_num_gradients_to_accumulate(new_accum)
-    for _ in range(new_accum):
-        out = model(Tensor([0.0, 1.0]))
-        out.sum().backward()
+    optim.set_scale(float(num_grads_to_accum))
+    optim.set_num_gradients_to_accumulate(num_grads_to_accum)
+    for _ in range(num_grads_to_accum):
+        loss = model(Tensor([0.0, 1.0])).sum() / num_grads_to_accum
+        if is_scaled_loss:
+            loss = loss / num_grads_to_accum
+        loss.backward()
+
+    if not is_scaled_loss:
+        optim.scale_grad_by_num_grads_to_accum()
 
     assert np.allclose(optim.gain(), exp_gain), optim.gain()
     optim.step()