remove dependencies on deepspeed and wandb

modeling_magma.py

@@ -24,7 +24,6 @@ import numpy as np
 import torch
 import torch.utils.checkpoint
 from torch import nn
-import wandb
 import torch.distributed as dist
 from transformers.modeling_utils import PreTrainedModel
 from transformers.activations import ACT2FN
@@ -282,12 +281,6 @@ class MagmaForCausalLM(MagmaPreTrainedModel):
         self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1
         self._padding_side = "left"  # set it to left by default, user can use setter to change padding_sides
 
-        try:
-            if dist.get_rank() == 0:
-                wandb.init(project=os.environ['WANDB_PROJECT'])
-        except:
-            pass
-
         self.post_init()
     
     # def from_pretrained(self, pretrained_model_name_or_path, *model_args, **kwargs):
@@ -325,40 +318,6 @@ class MagmaForCausalLM(MagmaPreTrainedModel):
 
     def tie_weights(self):
         return self.language_model.tie_weights()
-
-    def load_special_module_from_ckpt(self, ckpt_path, torch_dtype=None):
-        from deepspeed.runtime.zero import Init
-        from deepspeed import zero
-        # Defer initialization for ZeRO-3 compatibility
-        # with Init(data_parallel_group=None):
-        #     # Initialize the special module
-        #     self.vision_tower = MagmaImageTower(self.config.vision_config, require_pretrained=False)
-
-        # Load checkpoint weights into the special module
-        checkpoint = torch.load(ckpt_path, map_location='cpu')
-        state_dict = {k.replace('visual.', ''): v for k, v in checkpoint.items() if 'visual.' in k}
-
-        # Convert checkpoint weights to match model's parameter dtype
-        if torch_dtype is None:
-            model_dtype = next(self.vision_tower.clip_vision_model.parameters()).dtype
-            for k, v in state_dict.items():
-                state_dict[k] = v.to(model_dtype)
-        else:
-            for k, v in state_dict.items():
-                state_dict[k] = v.to(torch_dtype)
-
-        # Temporarily gather parameters for loading (if ZeRO-3 is active)
-        with zero.GatheredParameters(list(self.vision_tower.parameters()), modifier_rank=0):
-            # Load the state dictionary
-            self.vision_tower.clip_vision_model.load_state_dict(state_dict, strict=False)
-            # After loading, ensure the module is on the correct device
-            for param in self.vision_tower.parameters():
-                param.data = param.data.to(self.device).to(torch_dtype)
-
-        # import pdb; pdb.set_trace()
-        # If using a DeepSpeed engine, attach the updated module
-        if hasattr(self, "deepspeed_engine"):
-            self.deepspeed_engine.module = self
         
     def resize_token_embeddings(self, new_num_tokens: Optional[int] = None, pad_to_multiple_of=None) -> nn.Embedding:
         model_embeds = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
@@ -832,13 +791,6 @@ class MagmaForCausalLM(MagmaPreTrainedModel):
             # concatenate the action accuracy across all devices
             action_accuracy = torch.cat(action_accuracy_gather)            
             
-            if dist.get_rank() == 0:            
-                # remove zero values
-                if action_accuracy.mean() == 0:
-                    wandb.log({"action_accuracy": action_accuracy.mean().item()})
-                else:
-                    action_accuracy = action_accuracy[action_accuracy != 0]
-                    wandb.log({"action_accuracy": action_accuracy.mean().item()})
         else:
             logits = self.language_model.lm_head(hidden_states)
             logits = logits.float()