additions

lm-evaluation-harness/lm_eval/models/my_olmoe.py

@@ -0,0 +1,358 @@
+"""
+LM Evaluation Harness Wrapper for Modified MyOLMoE
+"""
+import torch
+from typing import List, Optional, Union, Dict, Any
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from lm_eval.api.model import LM
+from lm_eval.api.registry import register_model
+import numpy as np
+
+
+@register_model("myolmoe")
+class MyOLMoELM(LM):
+    """LM Evaluation Harness wrapper for MYOLMoE model."""
+    
+    def __init__(
+        self,
+        pretrained: str = None,
+        device: str = "cuda",
+        batch_size: int = 1,
+        max_length: int = 2048,
+        trust_remote_code: bool = False,
+        dtype: str = "float16",
+        parallelize: bool = False,
+        device_map: Optional[str] = None,
+        **kwargs
+    ):
+        super().__init__()
+        
+        # Initialize device and batch size
+        if device == "cuda" and not torch.cuda.is_available():
+            device = "cpu"
+        self._device = torch.device(device)
+        self._batch_size = batch_size
+        self._max_length = max_length
+        
+        # Set dtype
+        if dtype == "float16":
+            self._dtype = torch.float16
+        elif dtype == "bfloat16":
+            self._dtype = torch.bfloat16
+        else:
+            self._dtype = torch.float32
+        
+        # Load tokenizer and model
+        if pretrained:
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                pretrained, 
+                trust_remote_code=trust_remote_code,
+                padding_side="left"
+            )
+            
+            # Ensure pad token is set
+            if self.tokenizer.pad_token is None:
+                if self.tokenizer.eos_token is not None:
+                    self.tokenizer.pad_token = self.tokenizer.eos_token
+                else:
+                    self.tokenizer.add_special_tokens({'pad_token': '<pad>'})
+            
+            self.model = AutoModelForCausalLM.from_pretrained(
+                pretrained,
+                torch_dtype=self._dtype,
+                device_map=device_map if parallelize else None,
+                trust_remote_code=trust_remote_code,
+                **kwargs
+            )
+            
+            if not parallelize:
+                self.model = self.model.to(self._device)
+            
+            self.model.eval()
+        else:
+            raise ValueError("pretrained model path must be specified")
+    
+    @property
+    def eot_token_id(self):
+        """End of text token ID."""
+        return self.tokenizer.eos_token_id
+    
+    @property
+    def max_length(self):
+        """Maximum sequence length."""
+        return self._max_length
+    
+    @property
+    def max_gen_toks(self):
+        """Maximum number of tokens to generate."""
+        return 256
+    
+    @property
+    def batch_size(self):
+        """Batch size for evaluation."""
+        return self._batch_size
+    
+    @property
+    def device(self):
+        """Device used for evaluation."""
+        return self._device
+    
+    def tok_encode(self, string: str, add_special_tokens=True) -> List[int]:
+        """Encode a string to token IDs."""
+        return self.tokenizer.encode(string, add_special_tokens=add_special_tokens)
+    
+    def tok_decode(self, tokens: List[int]) -> str:
+        """Decode token IDs to string."""
+        return self.tokenizer.decode(tokens, skip_special_tokens=True)
+    
+    def loglikelihood(self, requests: List[tuple]) -> List[tuple]:
+        """
+        Compute log-likelihood for each request.
+        Each request is a tuple of (context, continuation).
+        """
+        results = []
+        
+        # Process requests in batches
+        for i in range(0, len(requests), self.batch_size):
+            batch = requests[i:i + self.batch_size]
+            batch_results = self._loglikelihood_batch(batch)
+            results.extend(batch_results)
+        
+        return results
+    
+    def _loglikelihood_batch(self, batch: List[tuple]) -> List[tuple]:
+        """Process a batch of loglikelihood requests."""
+        contexts, continuations = zip(*batch)
+        
+        # Encode full sequences (context + continuation)
+        full_sequences = [ctx + cont for ctx, cont in zip(contexts, continuations)]
+        full_encodings = [self.tok_encode(seq) for seq in full_sequences]
+        
+        # Encode contexts only
+        context_encodings = [self.tok_encode(ctx) for ctx in contexts]
+        
+        # Pad sequences to the same length
+        max_len = min(max(len(seq) for seq in full_encodings), self.max_length)
+        
+        input_ids = []
+        attention_masks = []
+        continuation_masks = []
+        
+        for full_seq, ctx_seq in zip(full_encodings, context_encodings):
+            # Truncate if necessary (keep the end)
+            if len(full_seq) > max_len:
+                full_seq = full_seq[-max_len:]
+                ctx_len = max(0, len(ctx_seq) - (len(full_encodings[0]) - max_len))
+            else:
+                ctx_len = len(ctx_seq)
+            
+            # Create padding
+            pad_length = max_len - len(full_seq)
+            padded_seq = [self.tokenizer.pad_token_id] * pad_length + full_seq
+            attention_mask = [0] * pad_length + [1] * len(full_seq)
+            
+            # Create mask for continuation tokens only
+            continuation_mask = [0] * max_len
+            continuation_start = pad_length + ctx_len
+            for j in range(continuation_start, max_len):
+                continuation_mask[j] = 1
+            
+            input_ids.append(padded_seq)
+            attention_masks.append(attention_mask)
+            continuation_masks.append(continuation_mask)
+        
+        # Convert to tensors
+        input_ids = torch.tensor(input_ids, device=self.device)
+        attention_masks = torch.tensor(attention_masks, device=self.device)
+        continuation_masks = torch.tensor(continuation_masks, device=self.device)
+        
+        # Forward pass
+        with torch.no_grad():
+            outputs = self.model(input_ids=input_ids, attention_mask=attention_masks)
+            logits = outputs.logits
+        
+        # Compute log-likelihoods
+        results = []
+        for i in range(len(batch)):
+            # Get logits for positions where we predict continuation tokens
+            # Shift logits and tokens for next-token prediction
+            shifted_logits = logits[i, :-1]  # Remove last position
+            shifted_tokens = input_ids[i, 1:]  # Remove first position
+            shifted_mask = continuation_masks[i][1:]  # Remove first position
+            
+            # Only consider continuation tokens
+            valid_positions = shifted_mask.bool()
+            if valid_positions.sum() == 0:
+                results.append((float('-inf'), False))
+                continue
+            
+            # Get log probabilities
+            log_probs = torch.log_softmax(shifted_logits, dim=-1)
+            token_log_probs = log_probs.gather(1, shifted_tokens.unsqueeze(1)).squeeze(1)
+            
+            # Sum only over continuation tokens
+            valid_log_probs = token_log_probs[valid_positions]
+            total_log_prob = valid_log_probs.sum().item()
+            
+            # For simplicity, assume greedy is True
+            is_greedy = True
+            
+            results.append((total_log_prob, is_greedy))
+        
+        return results
+    
+    def generate_until(self, requests: List[tuple]) -> List[str]:
+        """
+        Generate text until stopping criteria are met.
+        Each request is a tuple of (context, generation_kwargs).
+        """
+        results = []
+        
+        # Process requests in batches
+        for i in range(0, len(requests), self.batch_size):
+            batch = requests[i:i + self.batch_size]
+            batch_results = self._generate_until_batch(batch)
+            results.extend(batch_results)
+        
+        return results
+    
+    def _generate_until_batch(self, batch: List[tuple]) -> List[str]:
+        """Process a batch of generation requests."""
+        contexts = []
+        gen_kwargs_list = []
+        
+        for context, gen_kwargs in batch:
+            contexts.append(context)
+            gen_kwargs_list.append(gen_kwargs)
+        
+        # Encode contexts
+        context_encodings = [self.tok_encode(ctx) for ctx in contexts]
+        
+        # Pad contexts
+        max_ctx_len = min(max(len(seq) for seq in context_encodings), 
+                         self.max_length - self.max_gen_toks)
+        
+        input_ids = []
+        attention_masks = []
+        
+        for ctx_seq in context_encodings:
+            # Truncate if necessary (keep the end)
+            if len(ctx_seq) > max_ctx_len:
+                ctx_seq = ctx_seq[-max_ctx_len:]
+            
+            # Pad sequence
+            pad_length = max_ctx_len - len(ctx_seq)
+            padded_seq = [self.tokenizer.pad_token_id] * pad_length + ctx_seq
+            attention_mask = [0] * pad_length + [1] * len(ctx_seq)
+            
+            input_ids.append(padded_seq)
+            attention_masks.append(attention_mask)
+        
+        # Convert to tensors
+        input_ids = torch.tensor(input_ids, device=self.device)
+        attention_masks = torch.tensor(attention_masks, device=self.device)
+        
+        # Generate
+        with torch.no_grad():
+            # Use first gen_kwargs for simplicity (can be extended)
+            gen_kwargs = gen_kwargs_list[0] if gen_kwargs_list else {}
+            
+            # Set default generation parameters
+            generation_kwargs = {
+                'max_new_tokens': gen_kwargs.get('max_gen_toks', self.max_gen_toks),
+                'do_sample': gen_kwargs.get('do_sample', False),
+                'temperature': gen_kwargs.get('temperature', 1.0),
+                'top_p': gen_kwargs.get('top_p', 1.0),
+                'pad_token_id': self.tokenizer.pad_token_id,
+                'eos_token_id': self.tokenizer.eos_token_id,
+                'attention_mask': attention_masks,
+                'use_cache': True,
+            }
+            
+            generated = self.model.generate(
+                input_ids=input_ids,
+                **generation_kwargs
+            )
+        
+        # Decode generated text
+        results = []
+        for i, gen_seq in enumerate(generated):
+            # Get original context length (without padding)
+            original_ctx_len = len(context_encodings[i])
+            
+            # Extract only the newly generated tokens
+            if len(gen_seq) > len(input_ids[i]):
+                new_tokens = gen_seq[len(input_ids[i]):].tolist()
+            else:
+                new_tokens = []
+            
+            # Decode
+            if new_tokens:
+                generated_text = self.tok_decode(new_tokens)
+            else:
+                generated_text = ""
+            
+            # Apply stopping criteria if specified
+            if 'until' in gen_kwargs_list[i]:
+                stop_strings = gen_kwargs_list[i]['until']
+                if isinstance(stop_strings, str):
+                    stop_strings = [stop_strings]
+                
+                for stop_str in stop_strings:
+                    if stop_str in generated_text:
+                        generated_text = generated_text[:generated_text.index(stop_str)]
+                        break
+            
+            results.append(generated_text)
+        
+        return results
+    
+    def loglikelihood_rolling(self, requests: List[tuple]) -> List[float]:
+        """
+        Compute rolling log-likelihood for each request.
+        Each request is a tuple containing the text to evaluate.
+        """
+        results = []
+        
+        for request in requests:
+            text = request[0] if isinstance(request, tuple) else request
+            tokens = self.tok_encode(text)
+            
+            if len(tokens) <= 1:
+                results.append(0.0)
+                continue
+            
+            # Compute log-likelihood using sliding window approach
+            total_log_prob = 0.0
+            total_tokens = 0
+            
+            # Use sliding window for long sequences
+            window_size = min(self.max_length, len(tokens))
+            
+            for i in range(1, len(tokens)):
+                # Define the window
+                start_idx = max(0, i - window_size + 1)
+                end_idx = i + 1
+                
+                window_tokens = tokens[start_idx:end_idx]
+                input_ids = torch.tensor([window_tokens], device=self.device)
+                
+                with torch.no_grad():
+                    outputs = self.model(input_ids=input_ids)
+                    logits = outputs.logits
+                
+                # Get log probability for the target token
+                target_pos = len(window_tokens) - 1
+                target_token = window_tokens[target_pos]
+                
+                if target_pos > 0:  # Ensure we have a position to predict from
+                    token_logits = logits[0, target_pos - 1]
+                    log_prob = torch.log_softmax(token_logits, dim=-1)[target_token].item()
+                    total_log_prob += log_prob
+                    total_tokens += 1
+            
+            # Return mean log-likelihood per token
+            avg_log_prob = total_log_prob / total_tokens if total_tokens > 0 else 0.0
+            results.append(avg_log_prob)
+        
+        return results

myolmoe/modeling_myolmoe.py

@@ -0,0 +1,413 @@
+"""
+Modified OLMoE Model with Configurable Routing
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from typing import Optional, Tuple, Union, List
+import math
+
+
+class MyOLMoEConfig(PretrainedConfig):
+    """Configuration class for OLMoE model with configurable routing."""
+    
+    model_type = "myolmoe"
+    
+    def __init__(
+        self,
+        vocab_size=50304,
+        hidden_size=768,
+        intermediate_size=3072,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        num_key_value_heads=None,
+        hidden_act="swish",
+        max_position_embeddings=2048,
+        initializer_range=0.02,
+        rms_norm_eps=1e-5,
+        use_cache=True,
+        pad_token_id=None,
+        bos_token_id=1,
+        eos_token_id=2,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        # MoE specific parameters
+        num_experts=8,
+        num_experts_per_tok=2,
+        router_aux_loss_coef=0.001,
+        # Routing configuration
+        routing_type="dense",  # "dense", "sparse", "non_deterministic"
+        router_temperature=1.0,  # For non-deterministic routing
+        **kwargs
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads or num_attention_heads
+        self.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        
+        # MoE parameters
+        self.num_experts = num_experts
+        self.num_experts_per_tok = num_experts_per_tok
+        self.router_aux_loss_coef = router_aux_loss_coef
+        
+        # Routing configuration
+        self.routing_type = routing_type
+        self.router_temperature = router_temperature
+        
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs
+        )
+
+
+class MyOLMoERouter(nn.Module):
+    """Configurable router for OLMoE experts."""
+    
+    def __init__(self, config: MyOLMoEConfig):
+        super().__init__()
+        self.config = config
+        self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
+        
+    def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Route tokens to experts based on configuration.
+        
+        Returns:
+            router_logits: Routing logits/probabilities
+            router_probs: Expert selection probabilities
+        """
+        batch_size, seq_len, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        
+        # Compute router logits
+        router_logits = self.gate(hidden_states)
+        
+        if self.config.routing_type == "dense":
+            # Dense routing: use all experts with softmax weights
+            router_probs = F.softmax(router_logits, dim=-1)
+            
+        elif self.config.routing_type == "sparse":
+            # Sparse routing: select top-k experts
+            router_probs = F.softmax(router_logits, dim=-1)
+            topk_weights, topk_indices = torch.topk(
+                router_probs, self.config.num_experts_per_tok, dim=-1
+            )
+            # Zero out non-selected experts
+            mask = torch.zeros_like(router_probs)
+            mask.scatter_(-1, topk_indices, 1.0)
+            router_probs = router_probs * mask
+            # Renormalize
+            router_probs = router_probs / router_probs.sum(dim=-1, keepdim=True)
+            
+        elif self.config.routing_type == "non_deterministic":
+            # Only consider first half of experts for top-k selection
+            num_first_half = self.config.num_experts // 2
+            router_probs = F.softmax(router_logits, dim=-1)
+            
+            # Create mask for first half experts
+            mask = torch.zeros_like(router_probs)
+            mask[:, :num_first_half] = 1.0
+            
+            # Apply mask and renormalize probabilities
+            masked_probs = router_probs * mask
+            masked_probs = masked_probs / (masked_probs.sum(dim=-1, keepdim=True) + 1e-8)
+            
+            # Select top-k from first half
+            topk_weights, topk_indices = torch.topk(
+                masked_probs[:, :num_first_half],  # Only look at first half
+                min(self.config.num_experts_per_tok, num_first_half),  # Don't exceed available experts
+                dim=-1
+            )
+            
+            # Create final mask
+            final_mask = torch.zeros_like(router_probs)
+            final_mask.scatter_(-1, topk_indices, 1.0)
+            router_probs = router_probs * final_mask
+            router_probs = router_probs / (router_probs.sum(dim=-1, keepdim=True) + 1e-8)
+            
+        else:
+            raise ValueError(f"Unsupported routing type: {self.config.routing_type}")
+            
+        router_logits = router_logits.view(batch_size, seq_len, -1)
+        router_probs = router_probs.view(batch_size, seq_len, -1)
+        
+        return router_logits, router_probs
+
+
+class MyOLMoEExpert(nn.Module):
+    """Individual expert in the MoE layer."""
+    
+    def __init__(self, config: MyOLMoEConfig):
+        super().__init__()
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        self.act_fn = self._get_activation_fn(config.hidden_act)
+        
+    def _get_activation_fn(self, activation):
+        if activation == "swish" or activation == "silu":
+            return F.silu
+        elif activation == "relu":
+            return F.relu
+        elif activation == "gelu":
+            return F.gelu
+        else:
+            raise ValueError(f"Unsupported activation: {activation}")
+    
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class MyOLMoEMLP(nn.Module):
+    """MoE MLP layer with configurable routing."""
+    
+    def __init__(self, config: MyOLMoEConfig):
+        super().__init__()
+        self.config = config
+        self.router = MyOLMoERouter(config)
+        self.experts = nn.ModuleList([
+            MyOLMoEExpert(config) for _ in range(config.num_experts)
+        ])
+        
+    def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        batch_size, seq_len, hidden_dim = hidden_states.shape
+        hidden_states_flat = hidden_states.view(-1, hidden_dim)
+        
+        # Route to experts
+        router_logits, router_probs = self.router(hidden_states)
+        router_probs_flat = router_probs.view(-1, self.config.num_experts)
+        
+        # Process through experts
+        expert_outputs = []
+        for i, expert in enumerate(self.experts):
+            expert_output = expert(hidden_states_flat)
+            expert_outputs.append(expert_output)
+        
+        expert_outputs = torch.stack(expert_outputs, dim=-1)  # [batch*seq, hidden, num_experts]
+        
+        # Combine expert outputs
+        output = torch.sum(expert_outputs * router_probs_flat.unsqueeze(1), dim=-1)
+        output = output.view(batch_size, seq_len, hidden_dim)
+        
+        # Compute auxiliary loss
+        aux_loss = self._compute_aux_loss(router_probs_flat, router_logits.view(-1, self.config.num_experts))
+        
+        return output, aux_loss
+    
+    def _compute_aux_loss(self, router_probs, router_logits):
+        """Compute auxiliary loss for load balancing."""
+        if self.config.router_aux_loss_coef == 0:
+            return torch.tensor(0.0, device=router_probs.device)
+        
+        # Load balancing loss
+        num_tokens = router_probs.shape[0]
+        expert_usage = router_probs.sum(dim=0) / num_tokens  # Average usage per expert
+        aux_loss = self.config.num_experts * torch.sum(expert_usage * expert_usage)
+        
+        return self.config.router_aux_loss_coef * aux_loss
+
+
+class MyOLMoEDecoderLayer(nn.Module):
+    """Transformer decoder layer with MoE MLP."""
+    
+    def __init__(self, config: MyOLMoEConfig):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = nn.MultiheadAttention(
+            config.hidden_size, 
+            config.num_attention_heads,
+            dropout=0.0,
+            batch_first=True
+        )
+        self.mlp = MyOLMoEMLP(config)
+        self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+    
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        
+        # Self attention
+        attn_output, _ = self.self_attn(
+            hidden_states, hidden_states, hidden_states,
+            attn_mask=attention_mask
+        )
+        hidden_states = residual + attn_output
+        
+        # MLP
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        mlp_output, aux_loss = self.mlp(hidden_states)
+        hidden_states = residual + mlp_output
+        
+        return hidden_states, aux_loss
+
+
+class MyOLMoEModel(PreTrainedModel):
+    """OLMoE model with configurable routing."""
+    
+    config_class = MyOLMoEConfig
+    
+    def __init__(self, config: MyOLMoEConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList([
+            MyOLMoEDecoderLayer(config) for _ in range(config.num_hidden_layers)
+        ])
+        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        
+        self.gradient_checkpointing = False
+        self.post_init()
+    
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            batch_size, seq_length = input_ids.shape
+        elif inputs_embeds is not None:
+            batch_size, seq_length, _ = inputs_embeds.shape
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+        
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        
+        hidden_states = inputs_embeds
+        
+        all_hidden_states = () if output_hidden_states else None
+        total_aux_loss = 0.0
+        
+        for decoder_layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+            
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=attention_mask,
+            )
+            
+            hidden_states = layer_outputs[0]
+            total_aux_loss += layer_outputs[1]
+        
+        hidden_states = self.norm(hidden_states)
+        
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+        
+        if not return_dict:
+            return tuple(v for v in [hidden_states, all_hidden_states] if v is not None)
+        
+        return {
+            'last_hidden_state': hidden_states,
+            'hidden_states': all_hidden_states,
+            'aux_loss': total_aux_loss
+        }
+
+
+class MyOLMoEForCausalLM(PreTrainedModel):
+    """MyOLMoE model for causal language modeling."""
+    
+    config_class = MyOLMoEConfig
+    
+    def __init__(self, config):
+        print("⚡ Using CUSTOM MyOLMoE implementation!")  # Will show during loading
+        super().__init__(config)
+        self.model = MyOLMoEModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        
+        self.post_init()
+    
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+    
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+    
+    def get_output_embeddings(self):
+        return self.lm_head
+    
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+    
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+        )
+        
+        hidden_states = outputs['last_hidden_state']
+        logits = self.lm_head(hidden_states)
+        
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = nn.CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+            
+            # Add auxiliary loss
+            if 'aux_loss' in outputs:
+                loss += outputs['aux_loss']
+        
+        if not return_dict:
+            output = (logits,) + tuple(v for k, v in outputs.items() if k != 'last_hidden_state')
+            return (loss,) + output if loss is not None else output
+        
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.get('hidden_states'),
+        )
+
+
+# Register the model
+from transformers import AutoConfig, AutoModelForCausalLM
+AutoConfig.register("myolmoe", MyOLMoEConfig)
+AutoModelForCausalLM.register(MyOLMoEConfig, MyOLMoEForCausalLM)