fix: make linter happy

Signed-off-by: Oleg S <[email protected]>

Author: RobotSail

src/instructlab/training/config.py

@@ -30,8 +30,8 @@ class DeepSpeedOffloadStrategy(Enum):
 
 # public API
 class DistributedBackend(Enum):
-    FSDP: str = "fsdp"
-    DEEPSPEED: str = "deepspeed"
+    FSDP = "fsdp"
+    DEEPSPEED = "deepspeed"
 
 
 # public API
@@ -121,6 +121,7 @@ class DeepSpeedOptions(BaseModel):
     save_samples: int | None = None
 
 
+# public API
 class DistillationConfig(BaseModel):
     """
     Config to use when performing knowledge distillation during training.
@@ -192,7 +193,7 @@ class TrainingArgs(BaseModel):
     weight_decay: float = Field(0.0, ge=0.0)
 
     # settings for knowledge distillation
-    distillation_options: Optional[DistillationConfig] = None
+    distillation_options: DistillationConfig | None = None
     use_distillation: bool = False
 
     mock_data: Optional[bool] = False

src/instructlab/training/main_ds.py

@@ -483,7 +483,10 @@ def train(
             else:
                 loss = output.loss
 
-            assert loss is not None, "loss cannot be equal to None!"
+            if loss is None:
+                raise ValueError(
+                    "received a value of `None` for loss after calculations, this should not happen"
+                )
             log_loss = loss.detach().item()
 
             num_loss_counted_tokens, micro_batch_size, log_loss = map(
@@ -1086,7 +1089,7 @@ def run_training(torch_args: TorchrunArgs, train_args: TrainingArgs) -> None:
     parser.add_argument(
         "--weight_decay",
         type=float,
-        default=1e-3,
+        default=0,
         help="Weight decay rate for optimizers that support it.",
     )
     parser.add_argument(