Fix fp16 training

apps/accelerate/chatllama/chatllama/rlhf/actor.py

@@ -5,6 +5,7 @@
 from beartype.typing import Tuple
 from einops import rearrange
 from torch.utils.data import DataLoader, Dataset
+from torch.cuda.amp import GradScaler
 
 from chatllama.rlhf.base_model import BaseModel, BaseTrainer
 from chatllama.rlhf.config import ConfigActor
@@ -213,6 +214,12 @@ def __init__(self, config: ConfigActor) -> None:
         # HF accelerate
         self.setup_accelerate()
 
+        # define the scaler needed for vanilla pytorch with mixed precision
+        if (not self.accelerate_enable) and (not self.deepspeed_enable):
+            self.scaler = GradScaler()
+        else:
+            self.scaler = None
+
     def add_eos_token(
         self, tokens: torch.Tensor, mask: torch.Tensor
     ) -> Tuple[torch.Tensor, torch.Tensor]:
@@ -260,13 +267,15 @@ def train(
 
         # get config parameters
         if self.deepspeed_enable:
-            batch_size = self.train_dataloader.batch_size
+            # get batch size from deepspeed
+            batch_size = self.model_engine.train_batch_size()
         elif self.accelerate_enable:
             batch_size = (
                 self.config.batch_size * self.accelerator.num_processes
             )
         else:
             batch_size = self.config.batch_size
+
         epochs = self.config.epochs
         device = self.config.device
         checkpoint_steps = self.config.checkpoint_steps
@@ -334,17 +343,44 @@ def train(
                     attention_mask = attention_mask.to(device)
 
                 # forward pass
-                if self.config.deepspeed_enable:
+                if self.deepspeed_enable:
                     est_output = self.model_engine(
                         training_input, attention_mask
                     )
-                else:
+                elif self.accelerate_enable:
                     est_output = self.model(training_input, attention_mask)
+                else:
+                    with torch.autocast(
+                        device_type=self.config.device_type,
+                        dtype=torch.float16,
+                    ):
+                        est_output = self.model(training_input, attention_mask)
 
                 # compute loss
-                est_output = rearrange(est_output, "b s v -> (b s) v")
-                training_output = rearrange(training_output, "b s -> (b s)")
-                loss = self.loss_function(est_output, training_output)
+                if (not self.accelerate_enable) and (
+                    not self.deepspeed_enable
+                ):
+
+                    # vanilla pytorch use autocast
+                    with torch.autocast(
+                        device_type=self.config.device_type,
+                        dtype=torch.float16,
+                    ):
+                        est_output = rearrange(est_output, "b s v -> (b s) v")
+                        training_output = rearrange(
+                            training_output, "b s -> (b s)"
+                        )
+                        loss = self.loss_function(est_output, training_output)
+                else:
+
+                    # deepspeed and accelerate use defualt
+                    est_output = rearrange(est_output, "b s v -> (b s) v")
+                    training_output = rearrange(
+                        training_output, "b s -> (b s)"
+                    )
+                    loss = self.loss_function(est_output, training_output)
+
+                # save training stats
                 self.append_training_stats(training_loss=loss.item())
 
                 # backward pass
@@ -358,8 +394,9 @@ def train(
                     self.scheduler.step()
                 else:
                     self.optimizer.zero_grad()
-                    loss.backward()
-                    self.optimizer.step()
+                    self.scaler.scale(loss).backward()
+                    self.scaler.step(self.optimizer)
+                    self.scaler.update()
                     self.scheduler.step()
 
                 # print progress

apps/accelerate/chatllama/chatllama/rlhf/base_model.py

@@ -156,6 +156,7 @@ def __init__(self, config: ConfigType) -> None:
 
             # load the model from model_folder
             self.load()
+            my_logger.success("Model loaded")
 
         else:
             # ActorCritic initialization
@@ -481,20 +482,35 @@ def setup_deepspeed(
         # initialize deepspeed
         self.model_engine = None
         if self.deepspeed_enable is True:
-            (
-                self.model_engine,
-                self.optimizer,
-                self.train_dataloader,
-                self.scheduler,
-            ) = deepspeed.initialize(
-                args=None,
-                model=self.model,
-                model_parameters=self.model.parameters(),
-                optimizer=self.optimizer,
-                lr_scheduler=self.scheduler,
-                training_data=self.train_dataset,
-                config=self.deepspeed_config_path,
-            )
+            if isinstance(self.config, Config):
+                (
+                    self.model_engine,
+                    self.optimizer,
+                    self.train_dataloader,
+                    self.scheduler,
+                ) = deepspeed.initialize(
+                    args=None,
+                    model=self.model,
+                    model_parameters=self.model.parameters(),
+                    optimizer=self.optimizer,
+                    lr_scheduler=self.scheduler,
+                    training_data=self.train_dataset,
+                    config=self.deepspeed_config_path,
+                )
+            else:
+                (
+                    self.model_engine,
+                    self.optimizer,
+                    self.train_dataloader,
+                    self.scheduler,
+                ) = deepspeed.initialize(
+                    args=None,
+                    model=self.model,
+                    model_parameters=self.model.parameters(),
+                    lr_scheduler=self.scheduler,
+                    training_data=self.train_dataset,
+                    config=self.deepspeed_config_path,
+                )
             my_logger.info("Training with DeepSpeed")
 
     @beartype

apps/accelerate/chatllama/chatllama/rlhf/config.py

@@ -87,6 +87,7 @@ class ConfigReward:
     accelerate_enable: bool = False
 
     debug: bool = False
+    device_type: str = "cuda"
 
 
 # just for naming consistency
@@ -168,6 +169,7 @@ class ConfigActor:
     peft_config_path: str
     checkpoint_name: Optional[str] = None
     debug: bool = False
+    device_type: str = "cuda"
 
 
 @dataclass
@@ -231,6 +233,7 @@ class ConfigTrainer:
     accelerate_enable: bool
     checkpoint_name: Optional[str] = None
     debug: bool = False
+    device_type: str = "cuda"
 
 
 class Config:
@@ -270,9 +273,16 @@ def __init__(
     ) -> None:
 
         # if not specified use the device available
+        if device is not None:
+            if ":" in str(device):
+                device_type = str(device).split(":")[0]
+            else:
+                device_type = str(device)
+
         if device is None:
             if torch.cuda.is_available():
                 device = torch.device("cuda")
+                device_type = "cuda"
             else:
                 raise ValueError("No GPU available...")
             # print(f"Current device used :{str(device)}")
@@ -292,17 +302,21 @@ def __init__(
         # Trainer Config
         trainer_dict["device"] = device
         trainer_dict["debug"] = debug
+        trainer_dict["device_type"] = device_type
         self.trainer = ConfigTrainer(**trainer_dict)
         # Actor Config
         actor_dict["device"] = device
         actor_dict["debug"] = debug
+        actor_dict["device_type"] = device_type
         self.actor = ConfigActor(**actor_dict)
         # Critic Config
         critic_dict["device"] = device
         critic_dict["debug"] = debug
+        critic_dict["device_type"] = device_type
         self.critic = ConfigCritic(**critic_dict)
         self.critic.is_reward = False
         # Reward Config
         reward_dict["device"] = device
         reward_dict["debug"] = debug
+        reward_dict["device_type"] = device_type
         self.reward = ConfigReward(**reward_dict)

apps/accelerate/chatllama/chatllama/rlhf/reward.py

@@ -4,6 +4,7 @@
 import torch
 from beartype import beartype
 from torch.utils.data import Dataset, DataLoader
+from torch.cuda.amp import GradScaler
 
 from chatllama.rlhf.base_model import BaseModel, BaseTrainer
 from chatllama.rlhf.config import ConfigReward
@@ -150,21 +151,14 @@ def __init__(self, config: ConfigReward) -> None:
 
         # load the model
         self.model = RewardModel(config)
+
         self.accelerate_enable = self.model.accelerate_enable
         self.deepspeed_enable = self.model.deepspeed_enable
 
         # optimizer
-        if self.deepspeed_enable:
-            import deepspeed
-
-            deepspeed.ops.op_builder.CPUAdamBuilder().load()
-            self.optimizer = deepspeed.ops.adam.DeepSpeedCPUAdam(
-                self.model.parameters(), lr=config.lr
-            )
-        else:
-            self.optimizer = torch.optim.AdamW(
-                self.model.parameters(), lr=config.lr
-            )
+        self.optimizer = torch.optim.AdamW(
+            self.model.parameters(), lr=config.lr
+        )
 
         # loss function
         self.loss_function = torch.nn.MSELoss()
@@ -197,12 +191,21 @@ def __init__(self, config: ConfigReward) -> None:
             last_epoch=-1,
         )
 
+        # for scaling the gradients
+        self.scaler = None
+
         # deepspeed
         self.setup_deepspeed()
 
         # HF accelerate
         self.setup_accelerate()
 
+        # define the scaler needed for vanilla pytorch with mixed precision
+        if (not self.accelerate_enable) and (not self.deepspeed_enable):
+            self.scaler = GradScaler()
+        else:
+            self.scaler = None
+
     def train(
         self,
     ) -> None:
@@ -212,7 +215,8 @@ def train(
 
         # get config parameters
         if self.deepspeed_enable:
-            batch_size = self.train_dataloader.batch_size
+            # get batch size from deepspeed
+            batch_size = self.model_engine.train_batch_size()
         elif self.accelerate_enable:
             batch_size = (
                 self.config.batch_size * self.accelerator.num_processes
@@ -263,30 +267,43 @@ def train(
                             truncation=True,
                             padding=True,
                         )
-                    output = torch.as_tensor(
-                        score, dtype=torch.float32, device=device
-                    )
+                    output = torch.as_tensor(score, device=device)
 
                 # forward pass
                 if self.config.deepspeed_enable:
                     est_output = self.model_engine(
                         input_tokens["input_ids"].to(device),
                         input_tokens["attention_mask"].to(device),
                     )[:, -1]
+                elif self.accelerate_enable:
+                    est_output = self.model.module.get_reward(
+                        input_tokens["input_ids"].to(device),
+                        input_tokens["attention_mask"].to(device),
+                    )
                 else:
-                    if self.accelerate_enable:
-                        est_output = self.model.module.get_reward(
-                            input_tokens["input_ids"].to(device),
-                            input_tokens["attention_mask"].to(device),
-                        )
-                    else:
+                    with torch.autocast(
+                        device_type=self.config.device_type,
+                        dtype=torch.float16,
+                    ):
                         est_output = self.model.get_reward(
                             input_tokens["input_ids"].to(device),
                             input_tokens["attention_mask"].to(device),
                         )
 
                 # compute the loss
-                loss = self.loss_function(est_output, output)
+                if (not self.accelerate_enable) and (
+                    not self.deepspeed_enable
+                ):
+                    # if vanilla pytorch use autocast
+                    with torch.autocast(
+                        device_type=self.config.device_type,
+                        dtype=torch.float16,
+                    ):
+                        loss = self.loss_function(est_output, output)
+                else:
+                    # compute the loss normally
+                    loss = self.loss_function(est_output, output)
+
                 self.append_training_stats(training_loss=loss.item())
 
                 # backward pass
@@ -300,8 +317,9 @@ def train(
                     self.scheduler.step()
                 else:
                     self.optimizer.zero_grad()
-                    loss.backward()
-                    self.optimizer.step()
+                    self.scaler.scale(loss).backward()
+                    self.scaler.step(self.optimizer)
+                    self.scaler.update()
                     self.scheduler.step()
 
                 # print progress