deepspeed-chat: print mean stage1/2 loss periodically

applications/DeepSpeed-Chat/dschat/rlhf/ppo_trainer.py

@@ -7,15 +7,15 @@
 import time
 import deepspeed
 from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
-from deepspeed.accelerator import get_accelerator
 
 from dschat.utils.utils import print_rank_0
 
 
 def print_all_ranks(tag, value, rank):
     world_size = torch.distributed.get_world_size()
-    all_tensor = torch.zeros(world_size, dtype=torch.float32).to(
-        get_accelerator().current_device_name())
+    all_tensor = torch.zeros(world_size,
+                             dtype=torch.float32,
+                             device=value.device)
     all_tensor[rank] = value
     torch.distributed.all_reduce(all_tensor, op=torch.distributed.ReduceOp.SUM)
     print_rank_0(f'{tag} {all_tensor}', rank)
@@ -55,7 +55,8 @@ def __init__(self, rlhf_engine, args):
         self.end_of_conversation_token_id = self.tokenizer(
             args.end_of_conversation_token)['input_ids'][-1]
         self.z3_enabled = args.actor_zero_stage == 3
-        self.compute_fp32_loss = self.args.compute_fp32_loss
+        self.calculate_fp32_loss = (self.args.dtype
+                                    == "bf16") and self.args.bf16_to_fp32_loss
 
         # In case the generated experience is not valid (too short), we use the last valid
         # generated experience. Alternatively, we can skip the step (on all workers).
@@ -70,6 +71,8 @@ def __init__(self, rlhf_engine, args):
         self.gamma = 1.0
         self.lam = 0.95
         self.generate_time = 0.0
+        self.first_generate = True
+        self.actor_model_hpu_graph_wrapped_fwd_fn = None
 
     def _generate_sequence(self, prompts, mask, step):
 
@@ -84,13 +87,32 @@ def _generate_sequence(self, prompts, mask, step):
             kwargs = dict()
 
         with torch.no_grad():
-            seq = self.actor_model.module.generate(
-                prompts,
-                attention_mask=mask,
-                max_length=max_min_length,
-                pad_token_id=self.tokenizer.pad_token_id,
-                synced_gpus=self.z3_enabled,
-                **kwargs)
+            if is_hpu():
+                import habana_frameworks.torch.hpu as thpu
+            if is_hpu() and self.args.enable_hpu_graphs:
+                orig_actor_model_fwd_fn = self.actor_model.module.forward
+                if self.first_generate:
+                    self.actor_model.module.forward = thpu.wrap_in_hpu_graph_func(
+                        self.actor_model.module.forward)
+                    self.first_generate = False
+                else:
+                    self.actor_model.module.forward = self.actor_model_hpu_graph_wrapped_fwd_fn
+                seq = self.actor_model.module.generate(
+                    prompts,
+                    attention_mask=mask,
+                    max_length=max_min_length,
+                    min_length=max_min_length,
+                    lazy_mode=True)
+                self.actor_model_hpu_graph_wrapped_fwd_fn = self.actor_model.module.forward
+                self.actor_model.module.forward = orig_actor_model_fwd_fn
+            else:
+                seq = self.actor_model.module.generate(
+                    prompts,
+                    attention_mask=mask,
+                    max_length=max_min_length,
+                    pad_token_id=self.tokenizer.pad_token_id,
+                    synced_gpus=self.z3_enabled,
+                    **kwargs)
 
         # Filter out seq with no answers (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
@@ -149,19 +171,25 @@ def generate_experience(self, prompts, mask, step):
 
         pad_token_id = self.tokenizer.pad_token_id
         attention_mask = seq.not_equal(pad_token_id).long()
+
+        hpu_mark_step()
         with torch.no_grad():
             output = self.actor_model(seq, attention_mask=attention_mask)
+            hpu_mark_step()
             output_ref = self.ref_model(seq, attention_mask=attention_mask)
+            hpu_mark_step()
             reward_score = self.reward_model.forward_value(
                 seq, attention_mask,
                 prompt_length=self.prompt_length)['chosen_end_scores'].detach(
                 )
+            hpu_mark_step()
             values = self.critic_model.forward_value(
                 seq, attention_mask, return_value_only=True).detach()[:, :-1]
+            hpu_mark_step()
 
         logits = output.logits
         logits_ref = output_ref.logits
-        if self.compute_fp32_loss:
+        if self.calculate_fp32_loss:
             logits = logits.to(torch.float)
             logits_ref = logits_ref.to(torch.float)
 
@@ -221,25 +249,34 @@ def train_rlhf(self, inputs):
             advantages, returns = self.get_advantages_and_returns(
                 old_values, old_rewards, start)
 
+        hpu_mark_step()
         ### process the new outputs
         batch = {'input_ids': seq, "attention_mask": attention_mask}
         actor_prob = self.actor_model(**batch, use_cache=False).logits
+        hpu_mark_step()
         actor_log_prob = gather_log_probs(actor_prob[:, :-1, :], seq[:, 1:])
+        hpu_mark_step()
         actor_loss = self.actor_loss_fn(actor_log_prob[:, start:],
                                         log_probs[:, start:], advantages,
                                         action_mask[:, start:])
+        hpu_mark_step()
         self.actor_model.backward(actor_loss)
+        hpu_mark_step()
 
         if not self.args.align_overflow:
             self.actor_model.step()
+            hpu_mark_step()
 
         value = self.critic_model.forward_value(**batch,
                                                 return_value_only=True,
                                                 use_cache=False)[:, :-1]
+        hpu_mark_step()
         critic_loss = self.critic_loss_fn(value[:, start:], old_values[:,
                                                                        start:],
                                           returns, action_mask[:, start:])
+        hpu_mark_step()
         self.critic_model.backward(critic_loss)
+        hpu_mark_step()
 
         if self.args.align_overflow:
             actor_overflow = self.actor_model.optimizer.check_overflow(
@@ -263,8 +300,10 @@ def train_rlhf(self, inputs):
                     "OVERFLOW: actor and critic overflow, skipping both actor and critic steps",
                     rank)
             self.actor_model.step()
+            hpu_mark_step()
 
         self.critic_model.step()
+        hpu_mark_step()
 
         return actor_loss, critic_loss
 
@@ -296,7 +335,7 @@ def critic_loss_fn(self, values, old_values, returns, mask):
             old_values - self.cliprange_value,
             old_values + self.cliprange_value,
         )
-        if self.compute_fp32_loss:
+        if self.calculate_fp32_loss:
             values = values.float()
             values_clipped = values_clipped.float()
         vf_loss1 = (values - returns)**2

applications/DeepSpeed-Chat/dschat/rlhf/rlhf_engine.py

@@ -7,6 +7,7 @@
 import deepspeed
 from deepspeed.ops.adam import FusedAdam
 from deepspeed.ops.adam import DeepSpeedCPUAdam
+from deepspeed.accelerator import get_accelerator
 from transformers import AutoModelForCausalLM, get_scheduler
 
 from dschat.utils.ds_utils import get_train_ds_config, get_eval_ds_config
@@ -104,8 +105,23 @@ def _init_actor(self, actor_model_name_or_path):
                 actor_model = make_model_gradient_checkpointing_compatible(
                     actor_model)
 
+        # TODO SW-146776: remove this WA once SW-141762 is resolved
+        if is_hpu():
+            import habana_frameworks.torch.core as htcore
+            actor_model.to(dtype=torch.bfloat16,
+                           device=get_accelerator().device())
+
         # Optimizer
-        AdamOptimizer = DeepSpeedCPUAdam if self.args.offload else FusedAdam
+        if self.args.offload:
+            AdamOptimizer = DeepSpeedCPUAdam
+        elif self.args.no_fused_kernels or is_hpu():
+            AdamOptimizer = torch.optim.AdamW
+        else:
+            AdamOptimizer = FusedAdam
+        print_rank_0(
+            f'Using {AdamOptimizer.__name__} optimizer for actor model',
+            self.args.global_rank)
+
         optim_params = get_optimizer_grouped_parameters(
             actor_model, self.args.actor_weight_decay,
             self.args.actor_lora_learning_rate)
@@ -234,8 +250,23 @@ def _init_critic(self, critic_model_name_or_path):
                 critic_model = make_model_gradient_checkpointing_compatible(
                     critic_model)
 
+        # TODO SW-146776: remove this WA once SW-141762 is resolved
+        if is_hpu():
+            critic_model.to(dtype=torch.bfloat16,
+                            device=get_accelerator().device())
+
         # Optimizer
-        AdamOptimizer = DeepSpeedCPUAdam if self.args.offload else FusedAdam
+        # TODO SW-147425: change the file to use HPEX optimizer instead of AdamW on hpu
+        if self.args.offload:
+            AdamOptimizer = DeepSpeedCPUAdam
+        elif self.args.no_fused_kernels or is_hpu():
+            AdamOptimizer = torch.optim.AdamW
+        else:
+            AdamOptimizer = FusedAdam
+        print_rank_0(
+            f'Using {AdamOptimizer.__name__} optimizer for critic model',
+            self.args.global_rank)
+
         optim_params = get_optimizer_grouped_parameters(
             critic_model, self.args.critic_weight_decay,
             self.args.critic_lora_learning_rate)

applications/DeepSpeed-Chat/dschat/utils/data/data_utils.py

@@ -298,8 +298,10 @@ def create_prompt_dataset(local_rank,
     eval_fname = f"{output_path}/evaldata_{fname}.pt"
 
     cache_found = os.path.isfile(train_fname) and os.path.isfile(eval_fname)
-    buf_create_cache = torch.ByteTensor([not cache_found]).to(
-        get_accelerator().current_device_name())
+    device = torch.device(get_accelerator().device_name(
+        torch.distributed.get_rank()))
+    buf_create_cache = get_accelerator().ByteTensor([not cache_found],
+                                                    device=device)
     torch.distributed.all_reduce(buf_create_cache)
 
     if local_rank <= 0 and (buf_create_cache.item() != 0 or reload):

applications/DeepSpeed-Chat/dschat/utils/model/model_utils.py

@@ -82,6 +82,70 @@ def causal_lm_forward(
     model.forward = causal_lm_forward
 
 
+def configure_dropout(model_config, dropout):
+    if dropout is not None:
+        for key in ('dropout', 'attention_dropout', 'hidden_dropout',
+                    'activation_dropout'):
+            if hasattr(model_config, key):
+                print(f"Setting model_config.{key} to {dropout}")
+                setattr(model_config, key, dropout)
+
+
+def causal_lm_model_to_fp32_loss(model):
+    """ Convert CausalLM model to calculate loss in fp32 """
+
+    def causal_lm_forward(
+        input_ids=None,
+        past_key_values=None,
+        attention_mask=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        **deprecated_arguments,
+    ):
+        output = model.__original_forward__(
+            input_ids=input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            labels=None,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict)
+
+        return_dict = isinstance(output, dict)
+        lm_logits = output.logits if return_dict else output[0]
+        loss = None
+        if labels is not None:
+            # move labels to correct device to enable model parallelism
+            labels = labels.to(lm_logits.device)
+            # Shift so that tokens < n predict n
+            shift_logits = lm_logits[..., :-1, :].float().contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            batch_size, seq_length, vocab_size = shift_logits.shape
+            # Flatten the tokens
+            loss_fct = torch.nn.CrossEntropyLoss()
+            loss = loss_fct(
+                shift_logits.view(batch_size * seq_length, vocab_size),
+                shift_labels.view(batch_size * seq_length))
+
+        if not return_dict:
+            # re-pack output with fp32 loss
+            return ((loss, ) + output) if loss is not None else output
+
+        output.loss = loss
+        return output
+
+    model.__original_forward__ = model.forward
+    model.forward = causal_lm_forward
+
+
 def create_hf_model(model_class,
                     model_name_or_path,
                     tokenizer,
@@ -122,7 +186,8 @@ def create_critic_model(model_name_or_path,
                         rlhf_training=False,
                         dropout=None,
                         zero_stage=0,
-                        compute_fp32_loss=False):
+                        loss_to_fp32=False,
+                        optimized_reward_loss_calc=False):
     # OPT model family always put a padding token at the beginning of the sequence,
     # we did not see this in other models but not sure if it is a general rule
 
@@ -139,7 +204,8 @@ def create_critic_model(model_name_or_path,
         critic_model,
         tokenizer,
         num_padding_at_beginning=num_padding_at_beginning,
-        compute_fp32_loss=compute_fp32_loss)
+        loss_to_fp32=loss_to_fp32,
+        opt_loss_calc=optimized_reward_loss_calc)
 
     if rlhf_training:
         # load critic model from checkpoint