Adopt PreTrainedModelWrapper for Hugging Face models

[jon-tow]

Summary

This PR adds the following changes:

• Adopts a modified version of the PreTrainedModelWrapper as implemented in the trl package, here, to allow for flexible wrapping of Hugging Face (HF) models. This is useful for providing intuitive access to Hugging Face PreTrainedModel attributes such as push_to_hub and save_pretrained without having to access the underlying model.

• Adds ILQL save_pretrained support`

• Introduces the following for HF-based architectures:

  • Renames architectures to AutoModelFor... to match Hugging Face counterparts.
  • Removes base_model.transformer and base_model.lm_head references and instead extracts the final hidden states from all hidden states in the forward pass. This saves 2x memory on save to disk where previously a model's state_dict stored both .transformer and .lm_head as separate states with the underlying transformer.

• Moves modeling code out of the trainer dir into a separate models dir.

• Renames some utils.modeling HF attribute getters from causal_lm to decoder as many of these utils are helpful in general e.g. for encoder-decoder (seq2seq) models as well (T5Branch uses the same getter for accessing decoder parts of the transformer).

TODOs

• Convert test scripts to proper unit tests for save_pretrained and from_pretrained for ILQL and PPO
• Add from_config method to support custom model configs, e.g. as used in some examples

  trlx/examples/randomwalks/ilql_randomwalks.py

  Line 24 in b2ce1a4

  GPT2Config(n_layer=6, n_embd=144, vocab_size=23),

Reports

• SFT Sentiments: https://api.wandb.ai/links/jon-tow/8euree8w
• RandomWalks (PPO & ILQL): https://api.wandb.ai/links/jon-tow/gv1z3lzt
• Sentiments (PPO): https://api.wandb.ai/links/jon-tow/gqcgkjha
• Sentiments (ILQL): https://api.wandb.ai/links/jon-tow/pvcxmos2
• Summarize CNN/DailyMail (T5 PPO): https://api.wandb.ai/links/jon-tow/brordspn

[cat-state]

Thanks for this PR, I like the reorganization and changes overall!

[cat-state]

thanks! this LGTM