DeepSpeed assumes model returns just one variable: loss

[g-karthik]

https://github.com/microsoft/DeepSpeed/blob/4d735946b8f256bc80ba13e3530f85c91d041ff4/deepspeed/pt/deepspeed_light.py#L582-L606

As you can see in line 596 above in the forward() call for the DeepSpeedLight engine, self.module (which gets initialized earlier with the model passed to deepspeed.initialize() by the client) is assumed to be returning just one output: loss. However, as a model developer, I can have many outputs being returned by my model, including several different losses. An example of such a model would be the GPT2DoubleHeadsModel in Hugging Face's transformers repo, which returns two different losses, one for each head/task.

The consequence of this is that I won't be able to integrate DeepSpeed as-is to work with such models. Could you please make the necessary changes to be able to support this use-case?

I suspect what you need to do is:

1. Move lines 599-600 (which perform loss scaling based on gradient accumulation steps) into your implementation of backward().
2. Update line 596 to reflect the fact that self.module could return a generic tuple of outputs instead of a loss.

This, in turn, will allow the client to define a customized loss by leveraging all the outputs, and then call your implementation of backward() with this loss as input. And things should be fine because the loss scaling is still happening, only in backward() instead of forward(). Does this make sense?

[g-karthik]

@jeffra @ShadenSmith

[tjruwase]

Hi Karthik,

Thanks for your question. Can you provide some clarification on this issue? In fact, DeepSpeed does not assume that loss is a single value. If you look _scale_loss() in https://github.com/microsoft/DeepSpeed/blob/4d735946b8f256bc80ba13e3530f85c91d041ff4/deepspeed/pt/deepspeed_light.py#L567 you will see that it handles tuples and lists of loss.

So are you seeing an actual bug or crash in a model with a tuple or list of loss values?

[g-karthik]

@tjruwase Thanks for the info!

Yes, I just noticed that loss could be of varying types as per _scale_loss(). This is a bit confusing as a name though, I recommend renaming this variable to outputs.

Having said that, there is still the assumption built into this method that loss only contains losses that need scaling to account for gradient accumulation. But GPT2DoubleHeadsModel from Hugging Face returns not just losses, but also hidden states and attentions. The code inside _scale_loss() will end up scaling down the hidden states and attentions too, which I think we don't want happening. It will also end up scaling the logits returned by the model during validation/inference happening periodically during training.

Hence I think my original suggestion to move lines 599-600 into backward() will resolve this problem.

I haven't run my code yet, I am still integrating DeepSpeed to it and diving into the library as I do it.

[g-karthik]

@tjruwase did you get a chance to look at this?

[tjruwase]

Apologies for the delay in responding to this.

It is not that straight forward. We used to handle scaling in backward() as you suggested, until we encountered client code that was scaling after forward(), for tensorboard plotting, and passed the scaled loss to backward(), where we blindly scaled again. Basically, we can not control or detect what the client did with the return of forward() before forwarding to backward(). Thus, there are downsides to scaling in forward and backward().

The rationale for going with scaling in forward is that it potentially relieves the client from dealing gradient accumulation logic. This might become more obvious as you write your script. One thing you will notice is that with DeepSpeed you will be able to avoid having different logic for gradient accumulation boundaries.

Also two points regarding your concern about scaling inappropriate outputs, like logits during validation:

1. Since this scaling in forward() is documented, the client can easily reverse the scaling
2. DeepSpeed is not targeting validation/inference at this point in time, as discussed in the other issue.

Hope that helps. Thanks so much for your engagement.

[g-karthik]

It is not that straight forward. We used to handle scaling in backward() as you suggested, until we encountered client code that was scaling after forward(), for tensorboard plotting, and passed the scaled loss to backward(), where we blindly scaled again. Basically, we can not control or detect what the client did with the return of forward() before forwarding to backward(). Thus, there are downsides to scaling in forward and backward().

Irrespective of whether you do loss scaling in forward() or backward(), you have no control over what the client does. So this rationale for moving scaling from backward() (as you state you did originally) to forward() does not make sense to me.

The rationale for going with scaling in forward is that it potentially relieves the client from dealing gradient accumulation logic. This might become more obvious as you write your script.

1. Since this scaling in forward() is documented, the client can easily reverse the scaling

Potentially relieves the client from dealing with gradient accumulation logic? I'm confused now and this is a bit of a vague explanation.

Let's suppose that the client is absolutely not doing any form of gradient accumulation logic and intends to fully utilize DeepSpeed for gradient accumulation. So the code to process a given batch looks something like (based on the discussions we've had on this and other issues, but please correct the code if wrong):

        model_obj = model_engine if args.deepspeed else model
        model_outputs = model_obj(*batch)
        loss_a = model_outputs[0]
        loss_b = model_outputs[1]
        other_outputs = model_outputs[2:]  # not being used in this particular code example

        loss = loss_a * args.coef_a + loss_b * args.coef_b

        if args.deepspeed:
            model_obj.backward(loss)
            model_obj.step()
        else:
               # do something else with model_obj

Now if your _scale_loss() were called in forward():

• other_outputs = model_outputs[2:] will also be scaled - this could include hidden states and attentions during training, and logits during validation/inference.

I see you have acknowledged that DeepSpeed is not intended to be used during validation/inference, so let's forget about that. However, during training, you will still end up scaling hidden states and attentions that are part of other_outputs, which is very client-unfriendly. Your suggestion that the client reverse the scaling being done by DeepSpeed to the non-loss train-time outputs is, simply put, a hack.

However, if your _scale_loss() were called at the start in backward():

• None of the elements in model_outputs would have been scaled prematurely.
• When model_obj.backward(loss) gets called, the only thing getting scaled will be the loss the client passed to backward(). Now if the client makes a mistake and passes some bad input, that's a different story. But at least the mistake will then lie with the client!

Can you please explain to me what I am missing here? It would be good if you used code to explain your arguments for why you cannot move the scaling to backward().

[tjruwase]

You raise a number of interesting issues here, thanks for the thorough analysis. I am now persuaded about the problems of scaling in forward(). At the end of this reply, I have put a sample of client codes that we encountered early on that shaped our thinking. I feel the sample code is less relevant to the discussion given the points you have made, but I have included it since you asked for it.

The example you use to make your point is more general than I had previously thought of. So to summarize my understanding/agreement with you on the problems of _scale_loss in forward():

1. It provides different semantics from client.forward(). This could be a problem in some cases. Ideally, we want DeepSpeed to provide a faster but functionally equivalent forward().
2. It makes assumptions about the output of client forward(), i.e., that they are scalar/composite loss values. Indeed, output == loss is not generally true.
3. It makes assumptions about the appropriate scaling function, i.e.. normalization by grad_accumulation_steps. Your example involves an arbitrary scaling function.
4. It destroys other non-loss related states, and requiring client recovery is indeed a hack (and does require client to manipulate gradient accumulation steps logic after all :)).

Given all of the above, I agree that _scale_loss should move out of forward. But, now I am skeptical whether loss scaling belongs in backward(), but rather that client should be expected to pass appropriately scaled loss values to backward(). In the obvious case, loss scaling is a one-liner (we not saving client much work), and in more exotic cases backward() can't replicate it anyways.

Since this is a breaking change, I will loop in the rest of the team (@jeffra , @ShadenSmith , @samyam ) on how to proceed. Thanks for expending much time on this issue.

model_obj = model_engine if args.deepspeed else model
for step,batch in enumerate(...):
  ...
  loss = model_obj(*batch)
  loss /= gradient_accumulation_steps
  print(f'...loss")
  engine.backward(loss)
  if step % gradient_accumulation_steps == 0:
      if args.deepspeed:
        model_obj.step()
      else: 
        optimizer.step()
        optimizer.zero_grad()
  else:
   if args.deepspeed:
      model_obj.step()  

[g-karthik]

@tjruwase Glad you agree with the points I've raised now! :)

3. It makes assumptions about the appropriate scaling function, i.e.. normalization by grad_accumulation_steps. Your example involves an arbitrary scaling function.

But, now I am skeptical whether loss scaling belongs in backward(), but rather that client should be expected to pass appropriately scaled loss values to backward(). In the obvious case, loss scaling is a one-liner (we not saving client much work), and in more exotic cases backward() can't replicate it anyways.

Indeed, my example involves arbitrary scaling of the two losses returned (via args.coef_a and args.coef_b). But please note that this scaling is independent of gradient accumulation. This scaling is for multi-task learning, where we want to control the importance of each task to the overall loss.

The gradient accumulation scaling needs to happen on the overall loss separately. This could be either internal to DeepSpeed or externally handled by the client script when not using DeepSpeed.

So I think moving gradient accumulation-related loss scaling to backward() will be just fine and should not be a breaking change. I'd also like to hear what @samyam, @jeffra and @ShadenSmith think.

[ShadenSmith]

I agree that the existing behavior should be changed, and it sounds like moving the scaling outside of forward() is the correct thing to do in the long run.

I'm more inclined to move _scale_loss() inside of backward() than to ask users to perform the gradient accumulation scaling themselves. Otherwise it breaks existing codes that use DeepSpeed with gradient accumulation. It also feels wrong for DeepSpeed to provide most of the mechanisms of gradient accumulation but require users to remember to do the scaling themselves.

However, I'm a little concerned about backwards compatibility. While moving the scaling to the beginning of backward() should not affect our own training or logging, it does change the existing behavior of DeepSpeed's forward pass since we would now be returning a non-scaled loss. So I believe this would be a "breaking" change for code that uses the returned loss for more than just passing to backward(). For example, user code that logs loss will observe a change. Is that a serious issue? I can envision someone seeing an abnormally large loss value and thinking the training is broken, or they need to perform the loss scaling themselves, etc.

If we're not worried about that, then I think it's as simple as moving the _scale_loss() call to backward(), document/warn users, and then probably bump our version to v0.2.0 to signify the backwards-incompatible change.

If we're worried about that confusion, we could keep scaling in forward() and add the ability to opt out? Then no existing codes need to change, but more general situations are handled correctly (albeit with some additional configuration required to match semantics of non-DeepSpeed forward()). To make things easier for users, we could then add an optional user-facing method to explicitly do the loss scaling (e.g., model_engine.scale_loss(loss)). It's optional in the sense that backward() could detect whether it has already been scaled and do so if not. This gives users the ability to access the scaled loss between forward() and backward(), but does not require them to do so. In summary, we support doing the loss scaling in either forward(), backward(), or at some point between the two. We scale in forward() by default for backwards compatibility and maintain some internal state to ensure that loss is scaled exactly once.

So the example from this thread could look like below, and loss will automatically scaled at the beginning of backward():

model_outputs = model_engine(*batch, scale_loss=False) # note the extra kwarg here
loss_a = model_outputs[0]
loss_b = model_outputs[1]
other_outputs = model_outputs[2:]  # not being used in this particular code example

loss = loss_a * args.coef_a + loss_b * args.coef_b

model_engine.backward(loss)
model_engine.step()

The user could instead insert loss = model_engine.scale_loss(loss) before backward() to use the scaled loss themselves.

The downside of the above idea is that in cases like @g-karthik 's, DeepSpeed will still silently do the wrong thing until scale_loss=False is specified. For that reason, breaking compatibility and just moving _scale_loss() to backward() still feels like a good or best option to me?

[tjruwase]

Shaden thanks for your comments. Based on the discussions so far, how about the following changes to preserve backward compatibility and enable users to enforce correct behavior.

1. Similar to Shaden's proposal, but slight difference. forward takes an extra scale_loss arg, but that defaults to true to preserve backward compatibility
   model_outputs = model_engine(*batch, scale_loss=True)

2. backward takes an extra scale_loss arg that defaults to False
   model_engine.backward(loss, scale_loss=False)

3. Call _scale_loss() in forward() and backward() based on scale_loss arg

4. Maintain state in deepspeed to detect and warn about 0 or 2 calls to _scale_loss() for a forward/backward sequence

What do you guys think?

[g-karthik]

I agree with @ShadenSmith, and I think this is overkill. Maintaining backward compatibility is important and one should strive to maintain it as much as possible, but in this specific context, I think the value of doing so is minimal. It is purely limited to logging purposes.

Moreover, there will likely be several actual breaking changes you'll need to make in future, given the nature of software and this field in general. As long as you document your breaking changes in the release notes for each new version, you should be fine.

[tjruwase]

So it sounds like the consensus is

1. Remove _scale_loss() from forward() so deepspeed forward is semantically equal to client forward
2. Expose model_engine.scale_loss(loss) so client can obtain scaled loss
3. Call _scale_loss() in backward(), assuming client is passing unscaled loss values

Is this correct? I can start preparing the PR.

[g-karthik]

@tjruwase I think instead of 2, you could just change the signature of backward() to also return the loss. The returned loss would be scaled, and typically users log their loss after forward+backward anyway, not just forward.

[tjruwase]

@g-karthik I like that idea. Explicit is good.