fix(ppo_gpt): prevent position_ids being None

[li-plus]

This PR fixes the below error during PPO training, by generating the correct position_ids when it is None.

The PPO training was launched by:

accelerate launch --num_processes 7 --config_file ../../configs/accelerate/zero2-bf16.yaml ppo_hh.py

The below error was raised:

Traceback (most recent call last):
  File "/data00/home/lijiahao.plus/deepspeed/trlx/examples/hh/ppo_hh.py", line 227, in <module>
    main(hparams)
  File "/data00/home/lijiahao.plus/deepspeed/trlx/examples/hh/ppo_hh.py", line 216, in main
    trlx.train(
  File "/data00/home/lijiahao.plus/deepspeed/trlx/trlx/trlx.py", line 103, in train
    trainer.make_experience(config.method.num_rollouts)
  File "/data00/home/lijiahao.plus/deepspeed/trlx/trlx/trainer/accelerate_ppo_trainer.py", line 408, in make_experience
    ref_logits = self.model.forward_hydra(
  File "/data00/home/lijiahao.plus/deepspeed/trlx/trlx/models/modeling_ppo.py", line 387, in forward_hydra
    hydra_outputs = self.frozen_head(input_hidden_state, output_shape, **forward_kwargs)
  File "/data00/home/lijiahao.plus/miniconda3/envs/mlir/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data00/home/lijiahao.plus/deepspeed/trlx/trlx/models/modeling_ppo.py", line 515, in forward
    outputs = block(
  File "/data00/home/lijiahao.plus/miniconda3/envs/mlir/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data00/home/lijiahao.plus/miniconda3/envs/mlir/lib/python3.9/site-packages/transformers/models/gpt_neox/modeling_gpt_neox.py", line 320, in forward
    attention_layer_outputs = self.attention(
  File "/data00/home/lijiahao.plus/miniconda3/envs/mlir/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data00/home/lijiahao.plus/miniconda3/envs/mlir/lib/python3.9/site-packages/transformers/models/gpt_neox/modeling_gpt_neox.py", line 139, in forward
    query, key = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids)
  File "/data00/home/lijiahao.plus/miniconda3/envs/mlir/lib/python3.9/site-packages/transformers/models/gpt_neox/modeling_gpt_neox.py", line 278, in apply_rotary_pos_emb
    gather_indices = position_ids[:, None, :, None]  # [bs, 1, seq_len, 1]
TypeError: 'NoneType' object is not subscriptable

[maxreciprocate]

Hi @li-plus, thanks for fixing this! For the context, the newest transformers==4.28.1 fixed NeoX modeling's left-padding issue, however it made position_ids a required argument (it must not be None). I also resolved few failed tests which resulted from the fact that GPTModelBranch is common for both and gpt-neox and gpt2 type models and the latter's GPT2Block.forward doesn't have position_ids argument. Now tests pass with both transformers==4.27.1 and transformers==4.28.1

[maxreciprocate]

It fixes #457 on transformers==4.28.1 as well

[jovany-wang]

So should we enable transformers==4.28.1 ?

[olliestanley]

@jovany-wang It would be great if you could enable transformers==4.28.1 if this PR has fixed the underlying issue. >=4.28.1 is required to use LLaMA or any derived models for example, so the current pin can be quite inconvenient for some use cases.

[maxreciprocate]

@olliestanley This will be resolved shortly with #465

[olliestanley]

@reciprocated That's great, thank you for your fast response!

[maxreciprocate]

@olliestanley it is fixed now