settings for fp8 training

Author: neph1

config/config_categories.yaml

@@ -1,5 +1,6 @@
-Dataset: data_root, video_column, caption_column, id_token, video_resolution_buckets, caption_dropout_p
-Training: training_type, seed, mixed_precision, train_steps, rank, lora_alpha, target_modules, gradient_accumulation_steps, checkpointing_steps, checkpointing_limit, enable_slicing, enable_tiling, batch_size
+Dataset: data_root, video_column, caption_column, id_token, video_resolution_buckets, caption_dropout_p, precompute_conditions
+Training: training_type, seed, train_steps, rank, lora_alpha, target_modules, gradient_accumulation_steps, checkpointing_steps, checkpointing_limit, enable_slicing, enable_tiling, batch_size
 Optimizer: optimizer, lr, beta1, beta2, epsilon, weight_decay, max_grad_norm, lr_scheduler, lr_num_cycles, lr_warmup_steps
 Validation: validation_steps, validation_epochs, num_validation_videos, validation_prompts, validation_prompt_separator
-Accelerate: gpu_ids, nccl_timeout, gradient_checkpointing, allow_tf32, dataloader_num_workers, report_to, accelerate_config
+Accelerate: gpu_ids, nccl_timeout, gradient_checkpointing, allow_tf32, dataloader_num_workers, report_to, accelerate_config
+Model: model_name, pretrained_model_name_or_path, text_encoder_dtype, text_encoder_2_dtype, text_encoder_3_dtype, vae_dtype, layerwise_upcasting_modules, layerwise_upcasting_storage_dtype, layerwise_upcasting_granularity

config/config_template.yaml

@@ -18,13 +18,15 @@ gpu_ids: '0'
 gradient_accumulation_steps: 4
 gradient_checkpointing: true
 id_token: afkx
+layerwise_upcasting_modules: [none, transformer]
+layerwise_upcasting_granularity: [pytorch_layer, diffusers_layer]
+layerwise_upcasting_storage_dtype: [float8_e4m3fn, float8_e5m2]
 lora_alpha: 128
 lr: 0.0001
 lr_num_cycles: 1
 lr_scheduler: ['linear', 'cosine', 'cosine_with_restarts', 'polynomial', 'constant', 'constant_with_warmup']
 lr_warmup_steps: 400
 max_grad_norm: 1.0
-mixed_precision: [bf16, fp16, 'no']
 model_name: ltx_video
 nccl_timeout: 1800
 num_validation_videos: 0
@@ -37,14 +39,14 @@ rank: 128
 report_to: none
 seed: 42
 target_modules: to_q to_k to_v to_out.0
-text_encoder_dtype: [bf16, fp16, fp32]
-text_encoder_2_dtype: [bf16, fp16, fp32]
-text_encoder_3_dtype: [bf16, fp16, fp32]
+text_encoder_dtype: [bf16, fp16, fp32, fp8]
+text_encoder_2_dtype: [bf16, fp16, fp32, fp8]
+text_encoder_3_dtype: [bf16, fp16, fp32, fp8]
 tracker_name: finetrainers
 train_steps: 3000
 training_type: lora
 use_8bit_bnb: false
-vae_dtype: [bf16, fp16, fp32]
+vae_dtype: [bf16, fp16, fp32, fp8]
 validation_epochs: 0
 validation_prompt_separator: ':::'
 validation_prompts: ''

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "finetrainers-ui"
-version = "0.8.0"
+version = "0.10.0"
 dependencies = [
     "gradio",
     "torch>=2.4.1"

run_trainer.py

@@ -20,7 +20,16 @@ def run(self, config: Config, finetrainers_path: str, log_file: str):
 
         # Model arguments
         model_cmd = f"--model_name {config.get('model_name')} \
-        --pretrained_model_name_or_path {config.get('pretrained_model_name_or_path')}"
+        --pretrained_model_name_or_path {config.get('pretrained_model_name_or_path')} \
+        --text_encoder_dtype {config.get('text_encoder_dtype')} \
+        --text_encoder_2_dtype {config.get('text_encoder_2_dtype')} \
+        --text_encoder_3_dtype {config.get('text_encoder_3_dtype')} \
+        --vae_dtype {config.get('vae_dtype')} "
+        
+        if config.get('layerwise_upcasting_modules') != 'none':
+            model_cmd += f"--layerwise_upcasting_modules {config.get('layerwise_upcasting_modules')} \
+            --layerwise_upcasting_storage_dtype {config.get('layerwise_upcasting_storage_dtype')} \
+            --layerwise_upcasting_granularity {config.get('layerwise_upcasting_granularity')} "
 
         # Dataset arguments
         dataset_cmd = f"--data_root {config.get('data_root')} \
@@ -30,11 +39,7 @@ def run(self, config: Config, finetrainers_path: str, log_file: str):
         --video_resolution_buckets {config.get('video_resolution_buckets')} \
         --caption_dropout_p {config.get('caption_dropout_p')} \
         --caption_dropout_technique {config.get('caption_dropout_technique')} \
-        {'--precompute_conditions' if config.get('precompute_conditions') else ''} \
-        --text_encoder_dtype {config.get('text_encoder_dtype')} \
-        --text_encoder_2_dtype {config.get('text_encoder_2_dtype')} \
-        --text_encoder_3_dtype {config.get('text_encoder_3_dtype')} \
-        --vae_dtype {config.get('vae_dtype')} "
+        {'--precompute_conditions' if config.get('precompute_conditions') else ''} "
 
         # Dataloader arguments
         dataloader_cmd = f"--dataloader_num_workers {config.get('dataloader_num_workers')}"
@@ -45,7 +50,6 @@ def run(self, config: Config, finetrainers_path: str, log_file: str):
         # Training arguments
         training_cmd = f"--training_type {config.get('training_type')} \
         --seed {config.get('seed')} \
-        --mixed_precision {config.get('mixed_precision')} \
         --batch_size {config.get('batch_size')} \
         --train_steps {config.get('train_steps')} \
         --rank {config.get('rank')} \

tabs/general_tab.py

@@ -15,7 +15,7 @@ def __init__(self, title, config_file_path, allow_load=False):
 
         try:
             with self.settings_column:
-                inputs = self.update_form(self.config)
+                inputs = self.update_form()
                 self.components = OrderedDict(inputs)
                 children = []
                 for child in self.settings_column.children:

tabs/prepare_tab.py

@@ -20,7 +20,7 @@ def __init__(self, title, config_file_path, allow_load=False):
 
         try:
             with self.settings_column:
-                self.components = OrderedDict(self.update_form(self.config))
+                self.components = OrderedDict(self.update_form())
                 for i in range(len(self.settings_column.children)):
                     keys = list(self.components.keys())
                     properties[keys[i]] = self.settings_column.children[i]

tabs/tab.py

@@ -71,10 +71,10 @@ def add_buttons(self):
                 outputs=[self.save_status, self.config_file_box, *self.get_properties().values()]
             )
 
-    def update_form(self, config):
+    def update_form(self):
         inputs = dict()
 
-        for key, value in config.items():
+        for key, value in self.config.items():
             category = 'Other'
             for categories in self.config_categories.keys():
                 if key in self.config_categories[categories]:

tabs/training_tab.py

@@ -30,7 +30,7 @@ def __init__(self, title, config_file_path, allow_load=False):
 
         try:
             with self.settings_column:
-                inputs = self.update_form(self.config)
+                inputs = self.update_form()
                 self.components = OrderedDict(inputs)
                 children = []
                 for child in self.settings_column.children:

tabs/training_tab_legacy.py

@@ -17,7 +17,7 @@ def __init__(self, title, config_file_path, allow_load=False):
 
         try:
             with self.settings_column:
-                self.components = OrderedDict(self.update_form(self.config))
+                self.components = OrderedDict(self.update_form())
                 for i in range(len(self.settings_column.children)):
                     keys = list(self.components.keys())
                     properties[keys[i]] = self.settings_column.children[i]

trainer_config_validator.py

@@ -29,7 +29,6 @@ def validate(self):
             'lr_scheduler',
             'lr_warmup_steps',
             'max_grad_norm',
-            'mixed_precision',
             'model_name',
             'nccl_timeout',
             'optimizer',