PPO fp16

[maxreciprocate]

🐛 Describe the bug

Traceback (most recent call last):
  File "examples/ppo_sentiments.py", line 61, in <module>
    main()
  File "examples/ppo_sentiments.py", line 52, in main
    trlx.train(
  File "/trlx/trlx/trlx.py", line 81, in train
    orch.make_experience(config.method.num_rollouts)
  File "/trlx/trlx/orchestrator/ppo_orchestrator.py", line 161, in make_experience
    logits, *_, values = self.trainer.model(
  File "/trlx/.env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/trlx/.env/lib/python3.8/site-packages/deepspeed/utils/nvtx.py", line 11, in wrapped_fn
    return func(*args, **kwargs)
  File "/trlx/.env/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1836, in forward
    loss = self.module(*inputs, **kwargs)
  File "/trlx/.env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/trlx/trlx/trainer/nn/ppo_models.py", line 393, in forward
    transformer_outputs = self.base_model.transformer(**forward_kwargs)
  File "/trlx/.env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/trlx/.env/lib/python3.8/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 832, in forward
    inputs_embeds = self.wte(input_ids)
  File "/trlx/.env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/trlx/.env/lib/python3.8/site-packages/torch/nn/modules/sparse.py", line 160, in forward
    return F.embedding(
  File "/trlx/.env/lib/python3.8/site-packages/torch/nn/functional.py", line 2210, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: Expected tensor for argument #1 'indices' to have one of the following scalar types: Long, Int; but got torch.cuda.HalfTensor instead (while checking arguments for embedding)

compute_environment: LOCAL_MACHINE
deepspeed_config:
  deepspeed_multinode_launcher: standard
  gradient_accumulation_steps: 1
  gradient_clipping: 1.0
  offload_optimizer_device: none
  offload_param_device: none
  zero3_init_flag: false
  zero_stage: 2
distributed_type: DEEPSPEED
downcast_bf16: false
fsdp_config: {}
main_process_port: 1234
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 1
use_cpu: false

accelerate launch --num_processes 1 --config_file fp16-zero2-deepspeed.yaml examples/ppo_sentiments.py

Which trlX version are you using?

main

Additional system and package information

torch==1.13.0+cu116 deepspeed==0.8.0

[ZHAOTING]

I ran into the same problem and possibly found a way to get over it.

I am using deepspeed==0.7.3.
When the aforementioned RuntimeError happened, I noticed from its log that deepspeed config has the following output:

[2023-02-01 19:29:25,778] [INFO] [config.py:976:print_user_config]   json = {
    "train_batch_size": 16,
    "train_micro_batch_size_per_gpu": 4,
    "gradient_accumulation_steps": 4,
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu"
        },
        "offload_param": {
            "device": "cpu"
        },
        "stage3_gather_16bit_weights_on_model_save": false
    },
    "gradient_clipping": 1.0,
    "steps_per_print": inf,
    "fp16": {
        "enabled": true,
        "auto_cast": true
    },
    "bf16": {
        "enabled": false
    },
    "zero_allow_untested_optimizer": true
}

The "fp16/autocast": true looks related to where the downcast to HalfTensor happens (L1663-1664 in the following screenshot).

This DS config comes from the Accelerate config. Since I have no experience configuring DS options in an Accelerate config file, I defined a standalone DS config file, and let the Accelerate refer to the DS config file. Under this setting, the fp16/autocast option does not appear in the logged DS config (and its default value is False), and the HalfTensor problem also disappears.

The contents of the new Accelerate config and DS config are as follows.

Accelerate config:

compute_environment: LOCAL_MACHINE
deepspeed_config:
  deepspeed_config_file: ds_config_trlx.json
  zero3_init_flag: false
distributed_type: DEEPSPEED
downcast_bf16: 'no'
dynamo_backend: 'NO'
fsdp_config: {}
machine_rank: 0
main_training_function: main
megatron_lm_config: {}
num_machines: 1
num_processes: 1
rdzv_backend: static
same_network: true
use_cpu: false

DS config (i.e. ds_config_trlx.json)

{
  "train_micro_batch_size_per_gpu": 2,
  "gradient_accumulation_steps": 4,
  "fp16": {
    "enabled": true,
    "min_loss_scale": 0.5,
    "fp16_scale_tolerance": 0.25,
    "opt_level": "O2"
  },
  "zero_optimization": {
    "stage": 2,
    "offload_param": {
      "device": "cpu"
    },
    "offload_optimizer": {
      "device": "cpu"
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 5e8,
    "contiguous_gradients": true
  }
}