fixes

Author: RaymondLi0

fast_llm/layers/vision_encoder/adapter.py

@@ -18,18 +18,18 @@ class VisionAdapter(Layer):
 
     def __init__(self, config: VisionEncoderConfig, tensor_space: TensorSpace):
         super().__init__()
-        input_dim = tensor_space.get_tensor_dim(VisionEncoderDimNames.out_channels)
+        input_dim = tensor_space[VisionEncoderDimNames.out_channels]
         self._activation_type = config.adapter_activation_type
         self.layer_1 = Linear(
             input_dim,
-            tensor_space.get_tensor_dim(VisionEncoderDimNames.adapter_size),
+            tensor_space[VisionEncoderDimNames.adapter_size],
             bias=True,
             weight_init_method=init_normal_(std=config.adapter_init_method_std),
             bias_init_method=init_normal_(std=config.adapter_init_method_std),
         )
         self.layer_2 = Linear(
-            tensor_space.get_tensor_dim(VisionEncoderDimNames.adapter_size),
-            tensor_space.get_tensor_dim(TransformerDimNames.hidden),
+            tensor_space[VisionEncoderDimNames.adapter_size],
+            tensor_space[TransformerDimNames.hidden],
             bias=True,
             weight_init_method=init_normal_(std=config.adapter_init_method_std),
             bias_init_method=init_normal_(std=config.adapter_init_method_std),

fast_llm/layers/vision_encoder/patch_conv.py

@@ -19,23 +19,23 @@ def __init__(self, config: VisionEncoderConfig, tensor_space: TensorSpace):
         self._lr_scale = config.adapter_lr_scale
         self.weight = ParameterMeta.from_dims(
             (
-                self._tensor_space.get_tensor_dim(VisionEncoderDimNames.out_channels),
-                self._tensor_space.get_tensor_dim(VisionEncoderDimNames.in_channels),
-                self._tensor_space.get_tensor_dim(VisionEncoderDimNames.patch_size),
-                self._tensor_space.get_tensor_dim(VisionEncoderDimNames.patch_size),
+                self._tensor_space[VisionEncoderDimNames.out_channels],
+                self._tensor_space[VisionEncoderDimNames.in_channels],
+                self._tensor_space[VisionEncoderDimNames.patch_size],
+                self._tensor_space[VisionEncoderDimNames.patch_size],
             ),
             init_method=init_normal_(),
             lr_scale=self._lr_scale,
         )
         if config.conv_bias:
             self.bias = ParameterMeta.from_dims(
-                (self._tensor_space.get_tensor_dim(VisionEncoderDimNames.out_channels),),
+                (self._tensor_space[VisionEncoderDimNames.out_channels],),
                 init_method=init_normal_(),
-                lr_sclae=self._lr_scale,
+                lr_scale=self._lr_scale,
             )
         else:
             self.bias = None
-        self.norm = config.patch_norm.get_layer(tensor_space.get_tensor_dim(VisionEncoderDimNames.out_channels))
+        self.norm = config.patch_norm.get_layer(tensor_space[VisionEncoderDimNames.out_channels])
         self.stride = config.patch_size
 
     def forward(

fast_llm/models/gpt/model.py

@@ -173,12 +173,8 @@ def preprocess_meta(
                 VisionEncoderKwargs.image_std: image_std,
                 VisionEncoderKwargs.image_rescale_factor: image_rescale_factor,
                 VisionEncoderKwargs.rope_theta: self._config.vision_encoder.transformer.rotary.theta,
-                VisionEncoderKwargs.kv_channels: self._tensor_space.get_tensor_dim(
-                    VisionTransformerDimNames.kv_channels
-                ).size,
-                VisionEncoderKwargs.out_channels: self._tensor_space.get_tensor_dim(
-                    VisionEncoderDimNames.out_channels
-                ).size,
+                VisionEncoderKwargs.kv_channels: self._tensor_space[VisionTransformerDimNames.kv_channels].size,
+                VisionEncoderKwargs.out_channels: self._tensor_space[VisionEncoderDimNames.out_channels].size,
             }
         else:
             vision_kwargs = {}
@@ -226,7 +222,7 @@ def preprocess_meta(
             else (batch_dim, hidden_sequence_q_dim, hidden_dim)
         )
         if self._config.vision_encoder.enabled:
-            vision_hidden_dim = self._tensor_space.get_tensor_dim(VisionTransformerDimNames.hidden)
+            vision_hidden_dim = self._tensor_space[VisionTransformerDimNames.hidden]
             vision_hidden_dims = (
                 (hidden_sequence_q_dim, batch_dim, vision_hidden_dim)
                 if sequence_first

fast_llm/models/ssm/config.py

@@ -133,7 +133,6 @@ def get_handler_class(cls) -> type[CheckpointHandler]:
 
 
 class LlavaHybridHuggingfaceCheckpointFormat(GPTHuggingfaceCheckpointFormat):
-    support_optimizer: typing.ClassVar[bool] = False
     name: typing.ClassVar[str] = "llava_hybrid"
     vision_name: typing.ClassVar[str] = "pixtral"
     text_name: typing.ClassVar[str] = "apriel_ssm_thinker_hybrid"