support VLMs

app.py

@@ -12,24 +12,46 @@ def load_config_from_content(content):
         # Try parsing as JSON first
         try:
             config = json.loads(content)
-            # Convert JSON HF config format to our format
-            return {
-                'hidden_size': config['hidden_size'],
-                'num_layers': config['num_hidden_layers'],
-                'vocab_size': config['vocab_size'],
-                'intermediate_size': config['intermediate_size'],
-                'seq_len': 2048,  # Default value since not in config
-                'mbs': 1,        # Default value
-                'batch_accum': 1, # Default value
-                'tp': 1,         # Default value
-                'pp': 1,         # Default value
-                'dp': 1,         # Default value
-                'zero_stage': 0,  # Default value
-                'tie_word_embeddings': config.get('tie_word_embeddings', True),
-                'num_attention_heads': config['num_attention_heads'],
-                'num_key_value_heads': config.get('num_key_value_heads', config['num_attention_heads']),
-                'fsdp_checkpointing': False  # Default value
-            }
+            # Check if this is a multimodal model with text_config
+            if 'text_config' in config:
+                # Use text_config for model parameters
+                text_config = config['text_config']
+                return {
+                    'hidden_size': text_config['hidden_size'],
+                    'num_layers': text_config['num_hidden_layers'],
+                    'vocab_size': config.get('vocab_size', 256000),  # Default for multimodal models
+                    'intermediate_size': text_config['intermediate_size'],
+                    'seq_len': 2048,  # Default value since not in config
+                    'mbs': 1,        # Default value
+                    'batch_accum': 1, # Default value
+                    'tp': 1,         # Default value
+                    'pp': 1,         # Default value
+                    'dp': 1,         # Default value
+                    'zero_stage': 0,  # Default value
+                    'tie_word_embeddings': config.get('tie_word_embeddings', True),
+                    'num_attention_heads': text_config['num_attention_heads'],
+                    'num_key_value_heads': text_config.get('num_key_value_heads', text_config['num_attention_heads']),
+                    'full_checkpointing': False  # Default value
+                }
+            else:
+                # Original code for non-multimodal models
+                return {
+                    'hidden_size': config['hidden_size'],
+                    'num_layers': config['num_hidden_layers'],
+                    'vocab_size': config['vocab_size'],
+                    'intermediate_size': config['intermediate_size'],
+                    'seq_len': 2048,  # Default value since not in config
+                    'mbs': 1,        # Default value
+                    'batch_accum': 1, # Default value
+                    'tp': 1,         # Default value
+                    'pp': 1,         # Default value
+                    'dp': 1,         # Default value
+                    'zero_stage': 0,  # Default value
+                    'tie_word_embeddings': config.get('tie_word_embeddings', True),
+                    'num_attention_heads': config['num_attention_heads'],
+                    'num_key_value_heads': config.get('num_key_value_heads', config['num_attention_heads']),
+                    'full_checkpointing': False  # Default value
+                }
         except json.JSONDecodeError:
             # If not JSON, try YAML
             config = yaml.safe_load(content)
@@ -55,7 +77,7 @@ def load_config_from_content(content):
                 'tie_word_embeddings': model_config['tie_word_embeddings'],
                 'num_attention_heads': model_config['num_attention_heads'],
                 'num_key_value_heads': model_config.get('num_key_value_heads', model_config['num_attention_heads']),
-                'fsdp_checkpointing': optimizer.get('fsdp_checkpointing', False)  # Add FSDP checkpointing from config
+                'full_checkpointing': optimizer.get('full_checkpointing', False)  # Renamed from fsdp_checkpointing
             }
     except Exception as e:
         raise gr.Error(f"Error parsing configuration: {str(e)}")
@@ -92,7 +114,7 @@ def format_config_display(config):
             "seq_len", "mbs", "batch_accum"
         ],
         "Parallelism": [
-            "tp", "pp", "dp", "zero_stage", "fsdp_checkpointing"
+            "tp", "pp", "dp", "zero_stage", "full_checkpointing"
         ]
     }
 
@@ -154,7 +176,7 @@ with gr.Blocks() as demo:
                         pp = gr.Number(1, label="Pipeline Parallelism")
                         dp = gr.Number(1, label="Data Parallelism")
                     zero_stage = gr.Radio([0, 1, 2, 3], value=0, label="ZeRO Stage")
-                    fsdp_checkpointing = gr.Checkbox(False, label="FSDP Activation Checkpointing")
+                    full_checkpointing = gr.Checkbox(False, label="Full Activation Checkpointing")
 
                 manual_submit = gr.Button("Calculate Memory (Manual Input)")
         with gr.Column(scale=2):
@@ -171,7 +193,7 @@ with gr.Blocks() as demo:
             plot1, plot2, config_display, oom_display,
             hidden_size, num_attention_heads, num_key_value_heads, num_layers, 
             vocab_size, intermediate_size, seq_len, mbs, batch_accum,
-            tp, pp, dp, zero_stage, tie_word_embeddings, fsdp_checkpointing
+            tp, pp, dp, zero_stage, tie_word_embeddings, full_checkpointing
         ]
     )
 
@@ -202,7 +224,7 @@ with gr.Blocks() as demo:
             config['dp'],
             config['zero_stage'],
             config['tie_word_embeddings'],
-            config['fsdp_checkpointing']
+            config['full_checkpointing']
         ]
 
     # Handle manual input
@@ -222,7 +244,7 @@ with gr.Blocks() as demo:
             'tie_word_embeddings': args[13],
             'num_attention_heads': args[1],
             'num_key_value_heads': args[2],
-            'fsdp_checkpointing': args[14]  # Add FSDP checkpointing
+            'full_checkpointing': args[14]  # Renamed from fsdp_checkpointing
         }
         return process_yaml_and_update_ui(config)
 
@@ -231,7 +253,7 @@ with gr.Blocks() as demo:
         inputs=[
             hidden_size, num_attention_heads, num_key_value_heads, num_layers, vocab_size, intermediate_size,
             seq_len, mbs, batch_accum, tp, pp, dp, zero_stage,
-            tie_word_embeddings, fsdp_checkpointing  # Add FSDP checkpointing
+            tie_word_embeddings, full_checkpointing  # Renamed from fsdp_checkpointing
         ],
         outputs=[plot1, plot2, config_display, oom_display]
     )

utils.py

@@ -51,7 +51,7 @@ def get_num_hidden_layers_in_pp(hidden_size, num_layers, vocab_size, intermediat
 def calculate_memory_components(
     hidden_size, num_attention_heads, num_key_value_heads, num_layers, vocab_size, intermediate_size,
     seq_len, mbs, batch_accum, tp, pp, dp, zero_stage,
-    tie_word_embeddings, fsdp_checkpointing=False
+    tie_word_embeddings, full_checkpointing=False
 ):
     # Calculate base components first
     if pp == 1:
@@ -107,7 +107,7 @@ def calculate_memory_components(
         base_activs = num_layers * decoder_layer_mib + cast_to_fp32 + sharded_cross_entropy
         
         # Apply activation reduction for FSDP checkpointing in ZeRO-3
-        if zero_stage == 3 and fsdp_checkpointing:
+        if zero_stage == 3 and full_checkpointing:
             activs = base_activs / dp  # Activation memory is reduced by dp factor with checkpointing
         else:
             activs = base_activs
@@ -163,12 +163,12 @@ def calculate_memory_components(
 def plot_memory_breakdown(
     hidden_size, num_attention_heads, num_key_value_heads, num_layers, vocab_size, intermediate_size,
     seq_len, mbs, batch_accum, tp, pp, dp, zero_stage,
-    tie_word_embeddings, fsdp_checkpointing=False
+    tie_word_embeddings, full_checkpointing=False
 ):
     results = calculate_memory_components(
         hidden_size, num_attention_heads, num_key_value_heads, num_layers, vocab_size, intermediate_size,
         seq_len, mbs, batch_accum, tp, pp, dp, zero_stage,
-        tie_word_embeddings, fsdp_checkpointing
+        tie_word_embeddings, full_checkpointing
     )
     memory_usage_peak_tbi = results["Aggregates"]["Peak Memory (TBI)"]