ColossalAI/applications/Chat/coati/ray
Hongxin Liu b5f0566363
[chat] add distributed PPO trainer (#3740)
* Detached ppo (#9)

* run the base

* working on dist ppo

* sync

* detached trainer

* update detached trainer. no maker update function

* facing init problem

* 1 maker 1 trainer detached run. but no model update

* facing cuda problem

* fix save functions

* verified maker update

* nothing

* add ignore

* analyize loss issue

* remove some debug codes

* facing 2m1t stuck issue

* 2m1t verified

* do not use torchrun

* working on 2m2t

* working on 2m2t

* initialize strategy in ray actor env

* facing actor's init order issue

* facing ddp model update issue (need unwarp ddp)

* unwrap ddp actor

* checking 1m2t stuck problem

* nothing

* set timeout for trainer choosing. It solves the stuck problem!

* delete some debug output

* rename to sync with upstream

* rename to sync with upstream

* coati rename

* nothing

* I am going to detach the replaybuffer from trainer and make it a Ray Actor. Two benefits: 1. support TP trainer. 2. asynchronized buffer operations

* experience_maker_holder performs target-revolving _send_experience() instead of length comparison.

* move code to ray subfolder

* working on pipeline inference

* apply comments

* working on pipeline strategy. in progress.

* remove pipeline code. clean this branch

* update remote parameters by state_dict. no test

* nothing

* state_dict sharding transfer

* merge debug branch

* gemini _unwrap_model fix

* simplify code

* simplify code & fix LoRALinear AttributeError

* critic unwrapped state_dict

---------

Co-authored-by: csric <richcsr256@gmail.com>

* [chat] add perfomance evaluator and fix bugs (#10)

* [chat] add performance evaluator for ray

* [chat] refactor debug arg

* [chat] support hf config

* [chat] fix generation

* [chat] add 1mmt dummy example

* [chat] fix gemini ckpt

* split experience to send (#11)

Co-authored-by: csric <richcsr256@gmail.com>

* [chat] refactor trainer and maker (#12)

* [chat] refactor experience maker holder

* [chat] refactor model init

* [chat] refactor trainer args

* [chat] refactor model init

* [chat] refactor trainer

* [chat] refactor experience sending logic and training loop args (#13)

* [chat] refactor experience send logic

* [chat] refactor trainer

* [chat] refactor trainer

* [chat] refactor experience maker

* [chat] refactor pbar

* [chat] refactor example folder (#14)

* [chat] support quant (#15)

* [chat] add quant

* [chat] add quant example

* prompt example (#16)

* prompt example

* prompt load csv data

* remove legacy try

---------

Co-authored-by: csric <richcsr256@gmail.com>

* [chat] add mmmt dummy example and refactor experience sending (#17)

* [chat] add mmmt dummy example

* [chat] refactor naive strategy

* [chat] fix struck problem

* [chat] fix naive strategy

* [chat] optimize experience maker sending logic

* [chat] refactor sending assignment

* [chat] refactor performance evaluator (#18)

* Prompt Example & requires_grad state_dict & sharding state_dict (#19)

* prompt example

* prompt load csv data

* remove legacy try

* maker models require_grad set to False

* working on zero redundancy update

* mmmt_prompt example; naive strategy requires_grad state_dict & sharding; maker model requires_no_grad.

* remove legacy examples

* remove legacy examples

* remove replay buffer tp state. bad design

---------

Co-authored-by: csric <richcsr256@gmail.com>

* state_dict sending adapts to new unwrap function (#20)

* prompt example

* prompt load csv data

* remove legacy try

* maker models require_grad set to False

* working on zero redundancy update

* mmmt_prompt example; naive strategy requires_grad state_dict & sharding; maker model requires_no_grad.

* remove legacy examples

* remove legacy examples

* remove replay buffer tp state. bad design

* opt benchmark

* better script

* nothing

* [chat] strategy refactor unwrap model

* [chat] strategy refactor save model

* [chat] add docstr

* [chat] refactor trainer save model

* [chat] fix strategy typing

* [chat] refactor trainer save model

* [chat] update readme

* [chat] fix unit test

* working on lora reconstruction

* state_dict sending adapts to new unwrap function

* remove comments

---------

Co-authored-by: csric <richcsr256@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>

* [chat-ray] add readme (#21)

* add readme

* transparent graph

* add note background

---------

Co-authored-by: csric <richcsr256@gmail.com>

* [chat] get images from url (#22)

* Refactor/chat ray (#23)

* [chat] lora add todo

* [chat] remove unused pipeline strategy

* [chat] refactor example structure

* [chat] setup ci for ray

* [chat-ray] Support LoRA trainer. LoRA weights reconstruction. (#24)

* lora support prototype

* lora support

* 1mmt lora & remove useless code

---------

Co-authored-by: csric <richcsr256@gmail.com>

* [chat] fix test ci for ray

* [chat] fix test ci requirements for ray

* [chat] fix ray runtime env

* [chat] fix ray runtime env

* [chat] fix example ci docker args

* [chat] add debug info in trainer

* [chat] add nccl debug info

* [chat] skip ray test

* [doc] fix typo

---------

Co-authored-by: csric <59389055+CsRic@users.noreply.github.com>
Co-authored-by: csric <richcsr256@gmail.com>
2023-06-07 10:41:16 +08:00
..
callbacks [chat] add distributed PPO trainer (#3740) 2023-06-07 10:41:16 +08:00
README.md [chat] add distributed PPO trainer (#3740) 2023-06-07 10:41:16 +08:00
__init__.py [chat] add distributed PPO trainer (#3740) 2023-06-07 10:41:16 +08:00
detached_replay_buffer.py [chat] add distributed PPO trainer (#3740) 2023-06-07 10:41:16 +08:00
detached_trainer_base.py [chat] add distributed PPO trainer (#3740) 2023-06-07 10:41:16 +08:00
detached_trainer_ppo.py [chat] add distributed PPO trainer (#3740) 2023-06-07 10:41:16 +08:00
experience_maker_holder.py [chat] add distributed PPO trainer (#3740) 2023-06-07 10:41:16 +08:00
lora_constructor.py [chat] add distributed PPO trainer (#3740) 2023-06-07 10:41:16 +08:00
utils.py [chat] add distributed PPO trainer (#3740) 2023-06-07 10:41:16 +08:00

README.md

Distributed PPO Training on Stage 3

Detach Experience Makers and Trainers

We can completely separate the trainers and makers.

  • The experience maker performs inference, produces experience, and remotely delivers it to the trainer (1).
  • The trainer consumes experience to train models, and periodically transmits new model parameters to the maker (2.1, 2.2).
  • Using an experience buffer to overlap transmission and computing.

In this manner, each node will work continuously without model idle time, and different optimization strategies can be applied for inference and training to meet the needs of speed or storage. It is also helpful for scalability.

DetachedPPOTrainer and ExperienceMakerHolder are Ray Actors (distinguished from Actor Model), representing Trainer and Experience Maker on the graph above, respectively.

More about Ray Core

Usage

See examples at ColossalAI/application/Chat/examples/ray

Setup Makers

  • define makers' environment variables :

    env_info_makers = [{
        'local_rank': '0',
        'rank': str(rank),
        'world_size': str(num_makers),
        'master_port': maker_port,
        'master_addr': master_addr
    } for rank in range(num_makers)]
    
    
  • define maker models :

    def model_fn():
        actor = get_actor_from_args(...)
        critic = get_critic_from_args(...)
        reward_model = get_reward_model_from_args(...)
        initial_model = get_actor_from_args(...)
        return actor, critic, reward_model, initial_model
    
    
  • set experience_holder_refs :

    experience_holder_refs = [
        ExperienceMakerHolder.options(
            name=f"maker_{i}",
            num_gpus=1,
            max_concurrency=2
        ).remote(
            detached_trainer_name_list=[f"trainer_{x}" for x in target_trainers(...)],
            model_fn=model_fn,
            ...)
        for i, env_info_maker in enumerate(env_info_makers)
    ]
    

    The names in the detached_trainer_name_list refer to the target trainers that the maker should send experience to. We set a trainer's name the same as a maker, by .options(name="str"). See below.

Setup Trainers

  • define trainers' environment variables :

    env_info_trainers = [{
        'local_rank': '0',
        'rank': str(rank),
        'world_size': str(num_trainers),
        'master_port': trainer_port,
        'master_addr': master_addr
    } for rank in range(num_trainers)]
    
  • define trainer models :

    def trainer_model_fn():
        actor = get_actor_from_args(...)
        critic = get_critic_from_args(...)
        return actor, critic
    
  • set trainer_refs :

    trainer_refs = [
        DetachedPPOTrainer.options(
            name=f"trainer{i}",
            num_gpus=1,
            max_concurrency=2
        ).remote(
            experience_maker_holder_name_list=[f"maker{x}" for x in target_makers(...)],
            model_fn = trainer_model_fn(),
            ...)
        for i, env_info_trainer in enumerate(env_info_trainers)
    ]
    

    The names in experience_maker_holder_name_list refer to the target makers that the trainer should send updated models to. By setting detached_trainer_name_list and experience_maker_holder_name_list, we can customize the transmission graph.

Launch Jobs

  • define data_loader :

    def data_loader_fn():
        return = torch.utils.data.DataLoader(dataset=dataset)
    
    
  • launch makers :

    wait_tasks = []
    for experience_holder_ref in experience_holder_refs:
        wait_tasks.append(
            experience_holder_ref.workingloop.remote(data_loader_fn(),
                                                     num_steps=experience_steps))
    
    
  • launch trainers :

    for trainer_ref in trainer_refs:
        wait_tasks.append(trainer_ref.fit.remote(total_steps, update_steps, train_epochs))
    
  • wait for done :

    ray.get(wait_tasks)
    

Flexible Structure

We can deploy different strategies to makers and trainers. Here are some notions.

2 Makers 1 Trainer

2 Makers 2 Trainer

Maker Inference Quantization

Tensor Parallel

TODO

  • Support LoRA
  • Support TP & PP