README.md

# Import standard libraries
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from typing import Optional, Union, Tuple, Dict, Any
import math

# Import Hugging Face Transformers modules
from transformers import (
    AutoTokenizer,
    PreTrainedModel,
    PretrainedConfig,
    GenerationMixin,
    Trainer,
    TrainingArguments,
    DataCollatorForLanguageModeling,
    pipeline)
from transformers.utils.doc import add_start_docstrings_to_model_forward, replace_return_docstrings
from datasets import Dataset as HFDataset
from torch.utils.data import Dataset
from transformers.modeling_outputs import CausalLMOutputWithPast

_CONFIG_FOR_DOC = "TinyQwen3Config"

TINY_QWEN3_INPUTS_DOCSTRING = r"""
TinyQwen3ForCausalLM input.

Args:
    input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
        Indices of input sequence tokens in the vocabulary.
    attention_mask (`torch.FloatTensor`, optional):
        Mask to avoid performing attention on padding token indices.
    labels (`torch.LongTensor`, optional):
        Labels for computing the language modeling loss.
"""

# === Custom Multi-Head Attention to avoid SDPA warnings ===
class CustomMultiHeadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads, dropout=0.1):
        super().__init__()
        assert embed_dim % num_heads == 0

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.scale = self.head_dim ** -0.5

        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=False)
        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=False)
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False)

        self.dropout = nn.Dropout(dropout)

    def forward(self, x, attention_mask=None):
        batch_size, seq_len, embed_dim = x.size()

        # Linear projections
        q = self.q_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # Scaled dot-product attention
        scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale

        # Apply causal mask for autoregressive generation
        causal_mask = torch.triu(torch.ones(seq_len, seq_len, device=x.device), diagonal=1).bool()
        scores = scores.masked_fill(causal_mask.unsqueeze(0).unsqueeze(0), float('-inf'))

        # Apply attention mask if provided
        if attention_mask is not None:
            attention_mask = attention_mask.unsqueeze(1).unsqueeze(1)
            scores = scores.masked_fill(attention_mask == 0, float('-inf'))

        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # Apply attention to values
        out = torch.matmul(attn_weights, v)
        out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim)

        # Final projection
        out = self.out_proj(out)

        return out, attn_weights

# === Mixture of Experts Layer ===
class MoeLayer(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_experts=4, k=1):
        super(MoeLayer, self).__init__()
        self.num_experts = num_experts
        self.k = k
        self.gate = nn.Linear(input_dim, num_experts)
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(input_dim, hidden_dim),
                nn.GELU(),
                nn.Linear(hidden_dim, input_dim)
            ) for _ in range(num_experts)
        ])

    def forward(self, x):
        batch_size, seq_len, embed_dim = x.shape

        # Compute gate logits and select top-k experts
        gate_logits = self.gate(x)  # [batch_size, seq_len, num_experts]
        weights, indices = torch.topk(gate_logits, self.k, dim=-1)
        weights = torch.softmax(weights, dim=-1)  # [batch_size, seq_len, k]

        # Compute outputs from all experts
        expert_outputs = []
        for expert in self.experts:
            expert_outputs.append(expert(x))
        expert_outputs = torch.stack(expert_outputs, dim=-1)  # [batch_size, seq_len, embed_dim, num_experts]

        # Combine expert outputs
        combined_output = torch.zeros_like(x)
        for i in range(self.k):
            expert_idx = indices[..., i]  # [batch_size, seq_len]
            weight = weights[..., i]  # [batch_size, seq_len]

            # Gather outputs from selected experts
            selected_output = torch.gather(
                expert_outputs,
                -1,
                expert_idx.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, embed_dim, -1)
            ).squeeze(-1)

            combined_output += selected_output * weight.unsqueeze(-1)

        return combined_output


# === Tiny Transformer Block with MoE ===
class TinyMoETransformerBlock(nn.Module):
    def __init__(self, embed_dim, num_heads=2, num_experts=4, k=1):
        super(TinyMoETransformerBlock, self).__init__()
        self.attn = CustomMultiHeadAttention(embed_dim, num_heads)
        self.moe = MoeLayer(embed_dim, embed_dim * 2, num_experts=num_experts, k=k)
        self.norm1 = nn.LayerNorm(embed_dim)
        self.norm2 = nn.LayerNorm(embed_dim)

    def forward(self, x, attention_mask=None):
        attn_out, _ = self.attn(x, attention_mask)
        x = self.norm1(x + attn_out)
        moe_out = self.moe(x)
        x = self.norm2(x + moe_out)
        return x


# === TinyQwen3 Model Config and Architecture ===
class TinyQwen3Config(PretrainedConfig):
    model_type = "tiny_qwen3"

    def __init__(
        self,
        vocab_size=151936,  # Match Qwen3-0.6B tokenizer vocab size
        embed_dim=128,
        num_layers=3,
        num_heads=2,
        num_experts=4,
        k=1,
        max_position_embeddings=2048,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.embed_dim = embed_dim
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.num_experts = num_experts
        self.k = k
        self.max_position_embeddings = max_position_embeddings


class TinyQwen3Simulator(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.token_emb = nn.Embedding(config.vocab_size, config.embed_dim)
        self.pos_emb = nn.Parameter(torch.randn(1, config.max_position_embeddings, config.embed_dim))
        self.layers = nn.ModuleList([
            TinyMoETransformerBlock(config.embed_dim, config.num_heads, config.num_experts, config.k)
            for _ in range(config.num_layers)
        ])
        self.final_norm = nn.LayerNorm(config.embed_dim)

    def forward(self, input_ids, attention_mask=None):
        batch_size, seq_len = input_ids.size()

        # Clamp input_ids to valid range
        input_ids = torch.clamp(input_ids, 0, self.token_emb.num_embeddings - 1)

        # Ensure sequence length doesn't exceed position embeddings
        seq_len = min(seq_len, self.pos_emb.size(1))
        input_ids = input_ids[:, :seq_len]

        x = self.token_emb(input_ids) + self.pos_emb[:, :seq_len, :]

        for layer in self.layers:
            x = layer(x, attention_mask)
        x = self.final_norm(x)
        return x


class TinyQwen3ForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = TinyQwen3Config
    base_model_prefix = "model"
    main_input_name = "input_ids"

    def __init__(self, config):
        super().__init__(config)
        self.model = TinyQwen3Simulator(config)
        self.lm_head = nn.Linear(config.embed_dim, config.vocab_size, bias=False)
        self.post_init()

    def post_init(self):
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()

    def get_input_embeddings(self):
        return self.model.token_emb

    def set_input_embeddings(self, value):
        self.model.token_emb = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    @add_start_docstrings_to_model_forward(TINY_QWEN3_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        **kwargs
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        """
        Forward pass of the TinyQwen3 model for causal language modeling.

        Returns:
            CausalLMOutputWithPast: Model outputs including loss and logits.
        """
        # Get hidden states from the model
        hidden_states = self.model(input_ids, attention_mask)

        # Apply language modeling head to get logits
        logits = self.lm_head(hidden_states)

        loss = None
        if labels is not None:
            # Shift labels for next token prediction
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, self.config.vocab_size),
                shift_labels.view(-1),
                ignore_index=-100
            )

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
        )

    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
        if past_key_values is not None:
            input_ids = input_ids[:, -1:]
        return {"input_ids": input_ids}


# === Dataset: Use Tokenized Text ===
class TokenizedTextDataset(Dataset):
    def __init__(self, texts, tokenizer, max_length=128):
        self.texts = texts
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = self.texts[idx]
        encodings = self.tokenizer(
            text,
            truncation=True,
            padding="max_length",
            max_length=self.max_length,
            return_tensors="pt"
        )
        input_ids = encodings["input_ids"].squeeze(0)

        # Clamp token IDs to valid range to prevent CUDA errors
        input_ids = torch.clamp(input_ids, 0, self.tokenizer.vocab_size - 1)

        return {"input_ids": input_ids, "labels": input_ids.clone()}


# === Main Execution ===
if __name__ == "__main__":
    import os
    import warnings

    # Suppress the sliding window attention warning
    warnings.filterwarnings("ignore", message=".*Sliding Window Attention.*")

    os.environ["CUDA_LAUNCH_BLOCKING"] = "1"  # For better CUDA error tracing
    os.environ["CUDA_VISIBLE_DEVICES"] = ""  # Hide all CUDA devices

    # Force CPU execution to avoid CUDA issues during debugging
    device = torch.device("cpu")
    torch.cuda.is_available = lambda: False  # Force torch to think CUDA is not available

    # Load Qwen3-0.6B tokenizer
    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B", trust_remote_code=True)

    # Add padding token if it doesn't exist
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    print(f"Tokenizer vocab size: {tokenizer.vocab_size}")

    # Sample text for training
    sample_texts = [
        "Artificial intelligence is a wonderful field of study.",
        "Deep learning enables machines to learn from data.",
        "Transformers have revolutionized NLP.",
        "Mixture of Experts makes large models efficient.",
        "Qwen3 is a powerful language model."
    ]

    # Test tokenization first
    print("Testing tokenization...")
    for i, text in enumerate(sample_texts[:2]):
        tokens = tokenizer(text, return_tensors="pt")
        print(f"Text {i}: {text}")
        print(f"Tokens: {tokens['input_ids']}")
        print(f"Max token ID: {tokens['input_ids'].max().item()}")
        print()

    # Create dataset
    train_dataset = TokenizedTextDataset(sample_texts, tokenizer, max_length=64)

    # Initialize TinyQwen3 model with Qwen3 vocab size
    print("Initializing model...")
    config = TinyQwen3Config(
        vocab_size=tokenizer.vocab_size,
        embed_dim=128,
        num_layers=2,  # Reduced for debugging
        num_heads=2,
        num_experts=2,  # Reduced for debugging
        k=1,
        max_position_embeddings=64  # Reduced for debugging
    )
    model = TinyQwen3ForCausalLM(config).to(device)

    print(f"Model vocab size: {model.config.vocab_size}")
    print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")

    # Test forward pass
    print("Testing forward pass...")
    test_input = torch.randint(0, min(1000, tokenizer.vocab_size), (1, 10)).to(device)
    try:
        with torch.no_grad():
            output = model(test_input)
            print(f"Forward pass successful! Output shape: {output.logits.shape}")
    except Exception as e:
        print(f"Forward pass failed: {e}")
        exit(1)

    # Create a simple training loop instead of using Trainer to avoid CUDA issues
    print("Starting manual training loop...")
    model.train()
    optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)

    # Create a simple DataLoader
    from torch.utils.data import DataLoader
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
    train_dataloader = DataLoader(train_dataset, batch_size=1, shuffle=True, collate_fn=data_collator)

    for epoch in range(1):
        print(f"Epoch {epoch + 1}")
        total_loss = 0
        for step, batch in enumerate(train_dataloader):
            # Move batch to device (CPU)
            batch = {k: v.to(device) for k, v in batch.items()}

            # Forward pass
            outputs = model(**batch)
            loss = outputs.loss

            # Backward pass
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

            if step % 2 == 0:
                print(f"Step {step}, Loss: {loss.item():.4f}")

            if step >= 5:  # Train for just a few steps
                break

        print(f"Average loss: {total_loss / min(len(train_dataloader), 6):.4f}")

    print("Training completed successfully!")

    # Save model and tokenizer
    print("Saving model...")
    model.save_pretrained("./tiny_qwen3_model")
    tokenizer.save_pretrained("./tiny_qwen3_model")

    # Test inference
    print("Testing inference...")
    try:
        pipe = pipeline(
            "text-generation",
            model="./tiny_qwen3_model",
            tokenizer="./tiny_qwen3_model",
            trust_remote_code=True,
            device=-1  # Force CPU
        )
        result = pipe("Explain the concept", max_new_tokens=20, do_sample=False)
        print("Generated text:", result)
    except Exception as e:
        print(f"Inference failed: {e}")
        # Try direct model inference
        model.eval()
        test_text = "Explain the concept"
        inputs = tokenizer(test_text, return_tensors="pt").to(device)
        with torch.no_grad():
            outputs = model.generate(
                inputs.input_ids,
                max_new_tokens=10,
                do_sample=False,
                pad_token_id=tokenizer.pad_token_id
            )
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        print("Direct generation:", generated_text)

###### Outputs
<!-- Loading tokenizer...
Tokenizer vocab size: 151643
Testing tokenization...
Text 0: Artificial intelligence is a wonderful field of study.
Tokens: tensor([[ 9286, 16488, 11229,   374,   264, 11117,  2070,   315,  3920,    13]])
Max token ID: 16488

Text 1: Deep learning enables machines to learn from data.
Tokens: tensor([[33464,  6832, 20081, 12645,   311,  3960,   504,   821,    13]])
Max token ID: 33464

Initializing model...
Model vocab size: 151643
Model parameters: 39,225,348
Testing forward pass...
Forward pass successful! Output shape: torch.Size([1, 10, 151643])
Starting manual training loop...
Epoch 1
Step 0, Loss: 11.9908
Step 2, Loss: 11.9382
Step 4, Loss: 11.9401
Average loss: 11.9586
Training completed successfully!
Saving model...
Testing inference...
Inference failed: The checkpoint you are trying to load has model type `tiny_qwen3` but Transformers does not recognize this architecture. This could be because of an issue with the checkpoint, or because your version of Transformers is out of date.

You can update Transformers with the command `pip install --upgrade transformers`. If this does not work, and the checkpoint is very new, then there may not be a release version that supports this model yet. In this case, you can get the most up-to-date code by installing Transformers from source with the command `pip install git+https://github.com/huggingface/transformers.git`
Direct generation: Explain the concepteous莫名 kali handleyarQUEST EDUCantedAndWaitucas -->