[BLOOM] Clean modeling code

[thomasw21]

What does this PR do?

THERE'S A BREAKING CHANGE: past_key_values changes in terms of format, is that acceptable @sgugger @patrickvonplaten ?

Changes

Make a pass at cleaning up some code:

• convert causal_mask to bool tensor and use it via mask.
• simplify causal_mask creation function.
• revert back to baddbmm instead of bmm, it was unclear why this was changed, and training codebase uses baddbmm.
• switch back attention_scores upcasting to fp32 as it mimics the training procedure the closest.
• remove self.layer_number normalization. It was introduced in Meg-DS to prevent overflow when computing attention scores in float16, but it doesn't seem to impact the model at all here.
• remove a reshape in .split_head() which reduces memory footprint and increases throughput
• remove multiple reshapes when computing QKV with past values. Though this introduces a BREAKING CHANGE where past_key instead of being stored in [batch_size, num_heads, seq_length, head_dim] it's stored in [batch_size * num_heads, head_dim, seq_length].
• explicit all dimensions when computing .view, .reshape
• standardize namings for batch_size, seq_length etc ...
• type hint improvements

Speed-up in generation

Using the following script:

from transformers import AutoTokenizer, AutoModelForCausalLM
from timeit import timeit

def main():
    model_name = "bigscience/bloom-350m"
    max_length = 50
    batch_size = 16
    model = AutoModelForCausalLM.from_pretrained(model_name).cuda()
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    texts = ["Hello my name is"] * batch_size
    input_ids = tokenizer.batch_encode_plus(texts, return_tensors="pt").to("cuda")

    print(timeit(lambda: model.generate(**input_ids, max_length=max_length), number=100))


if __name__ == "__main__":
    main()

Results on A100:

This branch: 67.85658120410517
main: 78.72558180289343

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[thomasw21]

Ran quick test on the generation speed:

from transformers import AutoTokenizer, AutoModelForCausalLM
from timeit import timeit

def main():
    model_name = "bigscience/bloom-350m"
    max_length = 50
    model = AutoModelForCausalLM.from_pretrained(model_name).cuda()
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    text = "Hello my name is"
    input_ids = tokenizer.encode(text, return_tensors="pt").cuda()

    print(timeit(lambda: model.generate(input_ids, max_length=max_length), number=100))


if __name__ == "__main__":
    main()

The results are the following on A100:

This branch: 62.44300797022879
Main: 70.0128909018822

[michaelbenayoun]

Awesome!

[Muennighoff]

Can remove the beta if it's 1 by default anyways?

Suggested change

beta=self.beta,

[thomasw21]

I don't like relying on default :D

[Muennighoff]

Really nice job! Some small nits; Did you already try rerunning any of the eval with this branch, e.g. human eval?

[Muennighoff]

This change is because we did not use apply_query_key_layer_scaling for the training of BLOOM models?

[thomasw21]

No we did, but I actually don't think we need it. I think the trick was used to be able to train fp16 models, but since our fp16 models are small and the big ones are bf16 that trick shouldn't be use anymore IMO.

[Muennighoff]

But don't we still need it then, as the checkpoints are trained with layer scaling?

[thomasw21]

Why? So that you mimick perfectly the logits? Not sure this matters, and we should make sure to obtain the maximum precision we can get from the weights IMO. I'm running a bunch of greedy_search on 175b right now to assess if that matters.

[LysandreJik]

Agree with @sgugger regarding the shape of the past key values: fine for me for the shape to be changed. We did change the format for this output two years ago in #9391 and #9596, which could also have been considered breaking.

I'd like for this change to be contained to BLOOM (i.e., not ported to GPT-2) until @patrickvonplaten can comment as he was very involved in that reshape.

[thomasw21]

Thanks @LysandreJik ! Yeah this PR is only changing BLOOM modeling.

[thomasw21]

Okay tested 176b on lambada:

## This branch 
{
  "results": {
    "lambada": {
      "ppl": 3.9275610127374367,
      "ppl_stderr": 0.08456328894237608,
      "acc": 0.6716475839316903,
      "acc_stderr": 0.006542638265686496
    }
  },
  "versions": {
    "lambada": 0
  },
}

## 4.21.0
{
  "results": {
    "lambada": {
      "ppl": 3.933199250271968,
      "ppl_stderr": 0.08471163223917309,
      "acc": 0.6708713370851931,
      "acc_stderr": 0.006546580975553108
    }
  },
  "versions": {
    "lambada": 0
  },
}

I think this PR is ready to be checked.

[sgugger]

All look good to me, and multi-GPU tests pass with/without Accelerate. Thanks!

[ydshieh]

Just wanna say I especially love the change regarding the new way the attention mask is computed and used.
It is much cleaner and could potentially avoid a few issues if we use large negative values and summing different masks.

While working on a previous PR #17306, I have a discussion on Slack which suggested a similar approach.
However, there are other considerations to apply this change to the long existing models in transformers.

For the record (the discussion on Slack):

I am wondering if another approach might be better: what if we don't add large negative values as currently done, and no need to do sth. like (1.0 - extended_attention_mask) * -10000.0 etc. We just pass attention_mask as 0 and 1 , and use it as bool mask, like torch.where(attn_mask, attn_weights, -torch.inf)
before sending it to softmax.
(we already do this for causal_mask , but there still use -1e4) The disadvantage of this approach is it requires more changes than just using torch.finfo(tensor.dtype).min .

IMO this requires too many changes and could confuse users that are used to our way of changing the attention_mask into large negative values. I'd prefer torch.finfo(tensor.dtype).min . Also I think XLA likes add operation more that tf.where / jnp.where