Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Empty partition cache #360

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
Open

Empty partition cache #360

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
209465b
Empty partition cache
tjruwase Apr 19, 2023
bb60894
Validate empty_partition_cache
tjruwase Apr 21, 2023
ce049be
Empty cache after generate
tjruwase Apr 21, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 15 additions & 1 deletion applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/ppo_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,11 @@ def gather_log_probs(logits, labels):
return log_probs_labels.squeeze(-1)


def assert_empty_partition_cache(model, tag):
avail_or_inflight_params = [p.ds_id for p in model.parameters() if hasattr(p, 'ds_id') and p.ds_status != ZeroParamStatus.NOT_AVAILABLE]
assert len(avail_or_inflight_params) == 0, \
f'{tag} empty_partition_cache failed to evict all params: remaining = {avail_or_inflight_params}'

class DeepSpeedPPOTrainer():

def __init__(self, rlhf_engine, args):
Expand All @@ -68,12 +73,14 @@ def __init__(self, rlhf_engine, args):
def _generate_sequence(self, prompts):

max_min_length = self.max_answer_seq_len + prompts.shape[1]

with torch.no_grad():
seq = self.actor_model.module.generate(prompts,
max_length=max_min_length,
min_length=max_min_length)

self.actor_model.empty_partition_cache()
assert_empty_partition_cache(self.actor_model, 'actor_model after generate')

# Filter out seq with no asnwers (or very short). This happens when users directly use the pre-training ckpt without supervised finetuning
# NOTE: this will causes each GPU has different number of examples
batch_size = seq.shape[0]
Expand Down Expand Up @@ -171,6 +178,9 @@ def train_rlhf(self, inputs):
action_mask[:, start:])
self.actor_model.backward(actor_loss)
self.actor_model.step()
self.actor_model.empty_partition_cache()
assert_empty_partition_cache(self.actor_model, 'actor_model after rlhf step')

value = self.critic_model.forward_value(**batch,
return_value_only=True,
use_cache=False)[:, :-1]
Expand All @@ -179,6 +189,8 @@ def train_rlhf(self, inputs):
returns, action_mask[:, start:])
self.critic_model.backward(critic_loss)
self.critic_model.step()
self.critic_model.empty_partition_cache()
assert_empty_partition_cache(self.critic_model, 'critic_model after rlhf step')

return actor_loss, critic_loss

Expand Down Expand Up @@ -267,5 +279,7 @@ def train_unsupervised(self, inputs, unsup_coef):
loss = outputs.loss
self.actor_model.backward(unsup_coef * loss)
self.actor_model.step()
self.actor_model.empty_partition_cache()
assert_empty_partition_cache(self.actor_model, 'actor_model after unsuper_step')

return loss