LMFuser 0.0.2__tar.gz → 0.0.4__tar.gz

{lmfuser-0.0.2 → lmfuser-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: LMFuser
-Version: 0.0.2
+Version: 0.0.4
 Summary: The LMFuser training framework.
 Project-URL: Homepage, https://github.com/TYTTYTTYT/LMFuser
 Project-URL: Documentation, https://github.com/TYTTYTTYT/LMFuser

{lmfuser-0.0.2 → lmfuser-0.0.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "LMFuser"
-version = "0.0.2"
+version = "0.0.4"
 requires-python = ">= 3.11"
 description = "The LMFuser training framework."
 readme = "README.md"

{lmfuser-0.0.2 → lmfuser-0.0.4}/src/lmfuser/runners/ddp_runner.py

@@ -257,7 +257,8 @@ class DDPRunner(Runner[DDPRunnerConfig]):
         model = getattr(self, '_model', None)
         if model is None:
             model = self.load_model()
-            if self.config.model_precision == 'fp16':
+            if self.config.model_precision.value() == 'fp16':
+                logger.critical(f'casting model to fp16')
                 model = model.half()
             self._model = DDPWraper(model)
         return self._model
@@ -322,20 +323,17 @@ class DDPRunner(Runner[DDPRunnerConfig]):
             self.scheduler.load_state_dict(self._scheduler_states)
 
         if scaler is None:
-            if self.config.use_amp and get_world_size() > 1:
-                amp_selection = self.config.amp_precision.value()
-                if amp_selection == 'fp16':
-                    enable_scaler = True
-                elif amp_selection == 'bf16':
-                    enable_scaler = False
-                else:
-                    raise ValueError(f'Unknown amp precision "{amp_selection}"')
-                if OLD_GRADSCALER:
-                    self.scaler = GradScaler(enabled=enable_scaler)
-                else:
-                    self.scaler = GradScaler(device=get_default_device_type(), enabled=enable_scaler) # type: ignore
+            if self.config.use_amp.value() and self.config.amp_precision.value() == 'fp16':
+                enable_scaler = True
             else:
-                self.scaler = None
+                enable_scaler = False
+            if self.config.model_precision.value() == 'fp16':
+                enable_scaler = False
+
+            if OLD_GRADSCALER:
+                self.scaler = GradScaler(enabled=enable_scaler)
+            else:
+                self.scaler = GradScaler(device=get_default_device_type(), enabled=enable_scaler) # type: ignore
         else:
             self.scaler = scaler
 
@@ -394,11 +392,7 @@ class DDPRunner(Runner[DDPRunnerConfig]):
         for acc_idx in range(self.config._num_acc_steps):
             with ExitStack() as stack:
                 # check whether to use amp
-                if all([
-                    self.config.use_amp,
-                    self.config.model_precision == 'fp32',
-                    get_world_size() > 1
-                ]):
+                if self.config.use_amp.value():
                     amp_selection = self.config.amp_precision.value()
                     assert amp_selection is not None
                     stack.enter_context(autocast(
@@ -438,6 +432,7 @@ class DDPRunner(Runner[DDPRunnerConfig]):
                 self.scaler.scale(loss).backward()
             else:
                 loss.backward()
+
         running_loss = dist_avg(running_loss)
         self.step_log({f'{task.__class__.__name__}/train/loss': running_loss})
         self.step_log({'train/epoch': self.epoch})
@@ -453,6 +448,8 @@ class DDPRunner(Runner[DDPRunnerConfig]):
 
         grad_norm_clip_val = self.config.grad_norm_clip.value()
         if grad_norm_clip_val is not None:
+            if self.scaler is not None:
+                self.scaler.unscale_(self.optimizer)
             norm = clip_grad_norm_(
                 parameters=self.model.parameters(), 
                 max_norm=grad_norm_clip_val