Upload folder using huggingface_hub

checkpoints/sft/final_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dcbace558f0fea7a597de65aec13d6bfb2db23926c02178b638b75febe5bf7a2
+size 2578583807

configs/pretrain.yaml

@@ -0,0 +1,138 @@
+# TaoNet T2 Configuration for Pretraining
+# DeepSeek MLA + RoPE with Hybrid Muon+AdamW Optimizer
+# Full BF16 precision (no quantization)
+
+# ============================================================================
+# Model Architecture - TaoNet (DeepSeek MLA + RoPE)
+# ============================================================================
+model:
+  architecture_type: taonet
+  vocab_size: 8192
+  hidden_dim: 1024
+  num_layers: 16
+  num_heads: 8
+  max_seq_length: 1024
+  
+  # TaoNet-specific: Multi-head Latent Attention (MLA) compression
+  d_latent_kv: 768
+  
+  # RoPE (Rotary Position Embedding) dimension per head
+  # With hidden_dim=1024 and num_heads=8, head_dim = 128
+  d_rope: 128
+  
+  # Feed-forward intermediate dimension
+  hidden_dim_ff: 3072
+  
+  # Dropout rate (low for stability with large models)
+  dropout: 0.02
+  
+  # Grouped Query Attention (1 = standard MLA, >1 = GQA)
+  gqa_groups: 1
+  
+  # Optional: Use factorized embedding for parameter efficiency
+  # vocab (8192) → rank (96) → hidden (1024)
+  use_factorized_embedding: true
+  d_embed_rank: 96
+  
+  # Weight initialization standard deviation
+  init_std: 0.02
+
+# ============================================================================
+# Dataset Configuration - Local JSONL
+# ============================================================================
+dataset:
+  local: true
+  jsonl_path: /home/student/Data/TaoData/pretrain.jsonl
+  text_field: text
+  max_samples: 6700000
+  samples_per_chunk: 1000
+  
+  # Tokenizer configuration
+  tokenizer_type: sentencepiece
+  tokenizer_path: tokenizer/tokenizer.model
+  tokenizer_threads: 4
+
+# ============================================================================
+# Training Hyperparameters
+# ============================================================================
+batch_size: 8
+num_epochs: 1  # Set to 10 for full training
+gradient_accumulation_steps: 32  # Effective batch: 8 × 32 = 256
+
+# Maximum gradient norm for clipping (prevents ternary instability)
+max_grad_norm: 1.0
+
+# ============================================================================
+# Optimizer - Hybrid Muon + AdamW
+# ============================================================================
+# Strategy:
+# - Muon: For 2D Linear weight matrices (orthogonal/SVD-based optimization)
+#   - 2D weights: learning_rate (3e-3)
+# - AdamW: For 1D parameters (biases, norms, embeddings)
+#   - 1D params: adamw_lr (3e-4) = 1/10 × learning_rate
+
+optimizer:
+  optimizer_type: hybrid_muon_adamw
+  
+  # Learning rate for Muon (2D weight matrices)
+  learning_rate: 3e-3
+  
+  # Learning rate for AdamW (1D parameters)
+  # Typically 1/10 of learning_rate to prevent over-updating 1D params
+  adamw_lr: 3e-4
+  
+  # L2 regularization (weight decay)
+  weight_decay: 0.01
+  
+  # Adam betas
+  betas: [0.9, 0.999]
+  
+  # Epsilon for numerical stability
+  eps: 1e-8
+
+# ============================================================================
+# Learning Rate Scheduler - 3-Phase Cosine with Warmup
+# ============================================================================
+# Phases:
+# 1. Warmup: 0 → 1.0 (300 steps, ~1.4% of training)
+# 2. Steady: 1.0 (constant for 5% of training)
+# 3. Decay: 1.0 → 0.1 (cosine decay for remaining 95%)
+
+scheduler:
+  scheduler_type: cosineWarmup
+  warmup_steps: 300
+  warmup_ratio: 0.0  # Ignored if warmup_steps > 0
+  steady_ratio: 0.05  # 5% of total training steps at peak LR
+  min_lr_ratio: 0.1   # Decay to 10% of peak LR
+  num_cycles: 0.5  # For compatibility (not used in 3-phase schedule)
+
+# ============================================================================
+# Data Type and Device
+# ============================================================================
+dtype: bfloat16  # Use BF16 for better convergence with large models
+device: cuda     # Use GPU for training
+
+# ============================================================================
+# Checkpointing and Validation
+# ============================================================================
+checkpoint_dir: checkpoints/test
+save_every_steps: 81920
+save_best_model: true
+keep_last_n_checkpoints: 3
+
+# Validation
+eval_every_steps: 8192
+eval_samples: 8000
+
+# ============================================================================
+# Logging
+# ============================================================================
+log_every_steps: 50
+aim_repo: .aim
+
+# ============================================================================
+# Miscellaneous
+# ============================================================================
+seed: 42
+num_workers: 0
+pin_memory: true

configs/pretrain_gamma.yaml

@@ -0,0 +1,116 @@
+# GammaNet Configuration for Pretraining
+# Gamma Space Model blocks in the TaoTrain causal LM shell
+
+# ============================================================================
+# Model Architecture - GammaNet (Gamma Space Model)
+# ============================================================================
+model:
+  architecture_type: gamma_net
+  vocab_size: 8192
+  hidden_dim: 768
+  num_layers: 10
+  num_heads: 8
+  max_seq_length: 1024
+
+  # Shared LM shell feed-forward / embedding controls
+  hidden_dim_ff: 2048
+  dropout: 0.02
+  use_factorized_embedding: false
+  d_embed_rank: 96
+
+  # GammaSpaceBlock-specific settings
+  gamma_hidden_dim: 256
+  gamma_dt_min: 1e-3
+  gamma_dt_max: 1e-1
+  gamma_dt_init: 1e-2
+  gamma_discretization: bilinear
+  gamma_prenorm: true
+  gamma_residual_scale: 1.0
+  gamma_activation: gelu
+  gamma_gate: true
+  gamma_use_D: true
+  gamma_kernel_mode: auto
+  gamma_kernel_threshold: 64
+  gamma_use_output_linear: true
+  gamma_gate_bias: 2.0
+  gamma_input_gate: true
+  gamma_input_gate_bias: 2.0
+  gamma_layer_scale_init: 0.1
+
+  # Weight initialization standard deviation for TaoTrain LM shell
+  init_std: 0.02
+
+# ============================================================================
+# Dataset Configuration - Local JSONL
+# ============================================================================
+dataset:
+  local: true
+  jsonl_path: /home/student/Data/TaoData/pretrain.jsonl
+  text_field: text
+  max_samples: 1000000
+  samples_per_chunk: 1000
+
+  # Tokenizer configuration
+  tokenizer_type: sentencepiece
+  tokenizer_path: tokenizer/tokenizer.model
+  tokenizer_threads: 4
+
+# ============================================================================
+# Training Hyperparameters
+# ============================================================================
+batch_size: 32
+num_epochs: 2
+gradient_accumulation_steps: 8
+max_grad_norm: 1.0
+
+# ============================================================================
+# Optimizer - Hybrid Muon + AdamW
+# ============================================================================
+optimizer:
+  optimizer_type: hybrid_muon_adamw
+  learning_rate: 5e-3
+  adamw_lr: 5e-4
+  weight_decay: 0.01
+  betas: [0.9, 0.999]
+  eps: 1e-8
+
+# ============================================================================
+# Learning Rate Scheduler - 3-Phase Cosine with Warmup
+# ============================================================================
+scheduler:
+  scheduler_type: cosineWarmup
+  warmup_steps: 300
+  warmup_ratio: 0.0
+  steady_ratio: 0.05
+  min_lr_ratio: 0.1
+  num_cycles: 0.5
+
+# ============================================================================
+# Data Type and Device
+# ============================================================================
+dtype: bfloat16
+device: cuda
+
+# ============================================================================
+# Checkpointing and Validation
+# ============================================================================
+checkpoint_dir: checkpoints/pretrain_gamma
+save_every_steps: 81920
+save_best_model: true
+keep_last_n_checkpoints: 3
+
+eval_every_steps: 8192
+eval_samples: 8000
+
+# ============================================================================
+# Logging
+# ============================================================================
+log_every_steps: 50
+aim_repo: .aim
+
+# ============================================================================
+# Miscellaneous
+# ============================================================================
+seed: 42
+num_workers: 0
+pin_memory: true

configs/rl_dpo.yaml

@@ -0,0 +1,60 @@
+# Example configuration for RL training (RL stage assumes you have a reward model)
+
+model:
+  architecture_type: transformer
+  vocab_size: 50257
+  hidden_dim: 256
+  num_layers: 8
+  num_heads: 8
+  dropout: 0.1
+  max_seq_length: 512
+  init_std: 0.02
+
+dataset:
+  dataset_name: allenai/real_toxicity_prompts
+  split: train
+  prompt_column: text
+  max_samples: 2000
+  cache_dir: .cache/datasets
+  tokenizer_threads: 1  # Number of background threads for tokenization (1-32 recommended)
+
+batch_size: 4
+num_epochs: 1
+gradient_accumulation_steps: 8
+max_grad_norm: 0.5
+
+optimizer:
+  optimizer_type: adamw
+  learning_rate: 1e-5
+  weight_decay: 0.0
+
+scheduler:
+  scheduler_type: linearWarmup
+  warmup_steps: 50
+
+dtype: bfloat16
+device: cuda
+
+checkpoint_dir: checkpoints/rl
+save_every_steps: 100
+save_best_model: false
+keep_last_n_checkpoints: 2
+
+eval_every_steps: 100
+eval_samples: 100
+
+log_every_steps: 10
+aim_repo: .aim
+
+# RL-specific settings
+rl_method: ppo  # or "dpo"
+reward_model_path: checkpoints/reward_model.pt  # Path to your reward model
+ppo_epochs: 4
+ppo_clip_ratio: 0.2
+entropy_coeff: 0.01
+value_loss_coeff: 1.0
+generation_max_length: 256
+
+seed: 42
+num_workers: 0
+pin_memory: true

configs/sft.yaml

@@ -0,0 +1,92 @@
+# Example configuration for supervised fine-tuning
+# Uses TaoNet (MLA+RoPE) architecture loaded from pretrained checkpoint
+
+# ============================================================================
+# Model Architecture - TaoNet (DeepSeek MLA + RoPE)
+# ============================================================================
+model:
+  architecture_type: taonet
+  vocab_size: 8192
+  hidden_dim: 1024
+  num_layers: 16
+  num_heads: 8
+  max_seq_length: 1024
+  
+  # TaoNet-specific: Multi-head Latent Attention (MLA) compression
+  d_latent_kv: 768
+  
+  # RoPE (Rotary Position Embedding) dimension per head
+  # With hidden_dim=1024 and num_heads=8, head_dim = 128
+  d_rope: 128
+  
+  # Feed-forward intermediate dimension
+  hidden_dim_ff: 3072
+  
+  # Dropout rate (low for stability with large models)
+  dropout: 0.02
+  
+  # Grouped Query Attention (1 = standard MLA, >1 = GQA)
+  gqa_groups: 1
+  
+  # Optional: Use factorized embedding for parameter efficiency
+  # vocab (8192) → rank (96) → hidden (1024)
+  use_factorized_embedding: true
+  d_embed_rank: 96
+  
+  # Weight initialization standard deviation
+  init_std: 0.02
+
+dataset:
+  split: train
+  instruction_column: input
+  response_column: output
+  
+  local: true
+  jsonl_path: /home/student/Data/TaoData/sft.jsonl
+  samples_per_chunk: 1000
+  max_samples: 160000
+  cache_dir: .cache/datasets
+  instruction_template: "{instruction}\n{response}"
+
+  # Tokenizer configuration
+  tokenizer_type: sentencepiece
+  tokenizer_path: tokenizer/tokenizer.model
+  tokenizer_threads: 4
+
+# SFT-specific configuration (these fields are in SFTConfig)
+checkpoint_path: "checkpoints/yarn8k/best_model.pt"
+user_token: "<user>"
+assistant_token: "<assistant>"
+response_loss_only: true
+
+batch_size: 8
+num_epochs: 1
+gradient_accumulation_steps: 4
+max_grad_norm: 1.0
+
+optimizer:
+  optimizer_type: adamw
+  learning_rate: 5e-5  # Lower LR for fine-tuning (vs 5e-4 pretrain base, 5e-3 Muon)
+  weight_decay: 0.01
+
+scheduler:
+  scheduler_type: linearWarmup
+  warmup_steps: 500  # Less aggressive warmup for fine-tuning
+
+dtype: bfloat16
+device: cuda
+
+checkpoint_dir: checkpoints/sft
+save_every_steps: 81920
+save_best_model: true
+keep_last_n_checkpoints: 2
+
+eval_every_steps: 8192
+eval_samples: 200
+
+log_every_steps: 10
+aim_repo: .aim
+
+seed: 42
+num_workers: 0
+pin_memory: true

configs/tokenizer.yaml

@@ -0,0 +1,35 @@
+# Example configuration for training a SentencePiece tokenizer from JSONL data
+
+# Dataset source - JSONL file
+jsonl_path: /home/student/Data/TaoData/pretrain.jsonl
+text_field: text  # Field name in JSON for text data
+
+# Tokenizer training parameters
+vocab_size: 8192  # Keep aligned with pretrain/sft model vocab_size
+model_type: unigram  # SentencePiece model type: unigram, bpe, char, word
+character_coverage: 0.9995
+
+# Output configuration
+output_dir: tokenizer
+tokenizer_prefix: tokenizer
+
+# Custom special tokens
+# Built-in tokens are managed by SentencePiece and resolved at runtime.
+# Entries here are registered as user-defined symbols and should encode as
+# single tokens, but SentencePiece does not guarantee their exact IDs.
+# Note: Use \n for newline token, \t for tab, etc.
+special_tokens:
+  - "\n"        # Newline token - quoted to preserve literal \n in YAML
+  - <think>     # Special token for chain-of-thought reasoning
+  - <user>      # User message token
+  - <assistant> # Assistant message token
+  - <image>     # Image token for multimodal models
+
+# Data sampling (optional)
+# Set to a number to train on only the first N samples from the JSONL file
+# Useful for quick testing or sub-sampling large datasets
+# Omit or set to null to use entire file
+max_samples: 1000000
+
+# Optional metadata
+tokenizer_name: tokenizer

configs/vlm.yaml

@@ -0,0 +1,86 @@
+# Example configuration for multimodal vision-language connector training
+# Intended training path: pretrain.yaml -> sft.yaml -> vlm.yaml
+
+model:
+  architecture_type: multimodal_wrapper
+  llm_architecture_type: taonet
+  vocab_size: 8192
+  hidden_dim: 768
+  num_layers: 10
+  num_heads: 8
+  max_seq_length: 1024
+  d_latent_kv: 512
+  d_rope: 64
+  hidden_dim_ff: 2048
+  dropout: 0.02
+  gqa_groups: 1
+  use_factorized_embedding: false
+  d_embed_rank: 96
+  init_std: 0.02
+
+  vision_encoder_type: cnn
+  vision_output_dim: 256
+  image_size: 224
+  vision_prefix_tokens: 10
+  image_token: <image>
+  cnn_channels: [32, 64, 128]
+  cnn_kernel_size: 3
+
+dataset:
+  local: true
+  jsonl_path: /home/student/Data/TaoData/vision_pretrain.jsonl
+  image_path_column: image
+  image_path_aliases: [image, image_path, image_file, file_name]
+  caption_prompt: Describe the image.
+  samples_per_chunk: 1000
+  max_samples: 160000
+  cache_dir: .cache/datasets
+
+  tokenizer_type: sentencepiece
+  tokenizer_path: tokenizer/tokenizer.model
+  tokenizer_threads: 4
+
+checkpoint_path: checkpoints/sft/final_model.pt
+user_token: <user>
+assistant_token: <assistant>
+response_loss_only: true
+
+freeze_llm: true
+unfreeze_last_n_layers: 2
+vision_learning_rate: 1e-4
+llm_learning_rate: 5e-5
+vision_prefix_tokens: 10
+image_token: <image>
+image_size: 224
+
+batch_size: 8
+num_epochs: 1
+gradient_accumulation_steps: 4
+max_grad_norm: 1.0
+
+optimizer:
+  optimizer_type: adamw
+  learning_rate: 1e-4
+  weight_decay: 0.01
+
+scheduler:
+  scheduler_type: linearWarmup
+  warmup_steps: 500
+
+dtype: bfloat16
+device: cuda
+
+checkpoint_dir: checkpoints/vlm
+save_every_steps: 81920
+save_best_model: true
+keep_last_n_checkpoints: 2
+
+eval_every_steps: 8192
+eval_samples: 200
+
+log_every_steps: 10
+aim_repo: .aim
+
+seed: 42
+num_workers: 0
+pin_memory: true

configs/vlm_sft.yaml

@@ -0,0 +1,86 @@
+# Example configuration for end-to-end multimodal supervised fine-tuning
+# Intended training path: pretrain.yaml -> sft.yaml -> vlm.yaml -> vlm_sft.yaml
+
+model:
+  architecture_type: multimodal_wrapper
+  llm_architecture_type: taonet
+  vocab_size: 8192
+  hidden_dim: 768
+  num_layers: 10
+  num_heads: 8
+  max_seq_length: 1024
+  d_latent_kv: 512
+  d_rope: 64
+  hidden_dim_ff: 2048
+  dropout: 0.02
+  gqa_groups: 1
+  use_factorized_embedding: false
+  d_embed_rank: 96
+  init_std: 0.02
+
+  vision_encoder_type: cnn
+  vision_output_dim: 256
+  image_size: 224
+  vision_prefix_tokens: 10
+  image_token: <image>
+  cnn_channels: [32, 64, 128]
+  cnn_kernel_size: 3
+
+dataset:
+  local: true
+  jsonl_path: /home/student/Data/TaoData/vision_sft.jsonl
+  image_path_column: image
+  image_path_aliases: [images, image_path, image_file, file_name]
+  caption_prompt: Describe the image.
+  samples_per_chunk: 1000
+  max_samples: 160000
+  cache_dir: .cache/datasets
+
+  tokenizer_type: sentencepiece
+  tokenizer_path: tokenizer/tokenizer.model
+  tokenizer_threads: 4
+
+checkpoint_path: checkpoints/vlm/final_model.pt
+user_token: <user>
+assistant_token: <assistant>
+response_loss_only: true
+
+freeze_llm: false
+unfreeze_last_n_layers: 0
+vision_learning_rate: 5e-5
+llm_learning_rate: 5e-5
+vision_prefix_tokens: 10
+image_token: <image>
+image_size: 224
+
+batch_size: 8
+num_epochs: 1
+gradient_accumulation_steps: 4
+max_grad_norm: 1.0
+
+optimizer:
+  optimizer_type: adamw
+  learning_rate: 5e-5
+  weight_decay: 0.01
+
+scheduler:
+  scheduler_type: linearWarmup
+  warmup_steps: 500
+
+dtype: bfloat16
+device: cuda
+
+checkpoint_dir: checkpoints/vlm_sft
+save_every_steps: 81920
+save_best_model: true
+keep_last_n_checkpoints: 2
+
+eval_every_steps: 8192
+eval_samples: 200
+
+log_every_steps: 10
+aim_repo: .aim
+
+seed: 42
+num_workers: 0
+pin_memory: true

configs/yarn4k.yaml

@@ -0,0 +1,75 @@
+# TaoNet 200M Configuration for YaRN Continued Pretraining (1k -> 4k)
+
+model:
+  architecture_type: taonet
+  vocab_size: 8192
+  hidden_dim: 1024
+  num_layers: 16
+  num_heads: 8
+  max_seq_length: 4096
+
+  d_latent_kv: 768
+  d_rope: 128
+  hidden_dim_ff: 3072
+  dropout: 0.02
+  gqa_groups: 1
+  use_factorized_embedding: true
+  d_embed_rank: 96
+  init_std: 0.02
+
+  rope_scale: 40.0
+  yarn_enabled: true
+  yarn_original_max_seq_length: 1024
+  yarn_alpha: 1.0
+
+dataset:
+  local: true
+  jsonl_path: /home/student/Data/TaoData/pretrain.jsonl
+  text_field: text
+  max_samples: 300000
+  samples_per_chunk: 1000
+
+  tokenizer_type: sentencepiece
+  tokenizer_path: tokenizer/tokenizer.model
+  tokenizer_threads: 4
+
+checkpoint_path: checkpoints/pretrain/final_model.pt
+
+batch_size: 2
+num_epochs: 1
+gradient_accumulation_steps: 64
+max_grad_norm: 1.0
+
+optimizer:
+  optimizer_type: hybrid_muon_adamw
+  learning_rate: 1.5e-3
+  adamw_lr: 1.5e-4
+  weight_decay: 0.01
+  betas: [0.9, 0.999]
+  eps: 1e-8
+
+scheduler:
+  scheduler_type: cosineWarmup
+  warmup_steps: 300
+  warmup_ratio: 0.0
+  steady_ratio: 0.05
+  min_lr_ratio: 0.1
+  num_cycles: 0.5
+
+dtype: bfloat16
+device: cuda
+
+checkpoint_dir: checkpoints/yarn4k
+save_every_steps: 4096
+save_best_model: true
+keep_last_n_checkpoints: 3
+
+eval_every_steps: 4096
+eval_samples: 4000
+
+log_every_steps: 50
+aim_repo: .aim
+
+seed: 42
+num_workers: 0
+pin_memory: true

configs/yarn8k.yaml

@@ -0,0 +1,75 @@
+# TaoNet 200M Configuration for YaRN Continued Pretraining (4k -> 8k)
+
+model:
+  architecture_type: taonet
+  vocab_size: 8192
+  hidden_dim: 1024
+  num_layers: 16
+  num_heads: 8
+  max_seq_length: 8192
+
+  d_latent_kv: 768
+  d_rope: 128
+  hidden_dim_ff: 3072
+  dropout: 0.02
+  gqa_groups: 1
+  use_factorized_embedding: true
+  d_embed_rank: 96
+  init_std: 0.02
+
+  rope_scale: 40.0
+  yarn_enabled: true
+  yarn_original_max_seq_length: 4096
+  yarn_alpha: 1.0
+
+dataset:
+  local: true
+  jsonl_path: /home/student/Data/TaoData/pretrain.jsonl
+  text_field: text
+  max_samples: 800000
+  samples_per_chunk: 1000
+
+  tokenizer_type: sentencepiece
+  tokenizer_path: tokenizer/tokenizer.model
+  tokenizer_threads: 4
+
+checkpoint_path: checkpoints/yarn4k/best_model.pt
+
+batch_size: 2
+num_epochs: 1
+gradient_accumulation_steps: 32
+max_grad_norm: 1.0
+
+optimizer:
+  optimizer_type: hybrid_muon_adamw
+  learning_rate: 7.5e-4
+  adamw_lr: 7.5e-5
+  weight_decay: 0.01
+  betas: [0.9, 0.999]
+  eps: 1e-8
+
+scheduler:
+  scheduler_type: cosineWarmup
+  warmup_steps: 300
+  warmup_ratio: 0.0
+  steady_ratio: 0.05
+  min_lr_ratio: 0.1
+  num_cycles: 0.5
+
+dtype: bfloat16
+device: cuda
+
+checkpoint_dir: checkpoints/yarn8k
+save_every_steps: 2048
+save_best_model: true
+keep_last_n_checkpoints: 3
+
+eval_every_steps: 2048
+eval_samples: 2000
+
+log_every_steps: 50
+aim_repo: .aim
+
+seed: 42
+num_workers: 0
+pin_memory: true

configs/yarn_pretrain.yaml

@@ -0,0 +1,181 @@
+# TaoNet T2 Configuration for YaRN Continued Pretraining
+# Extended Context: 1024 → 8192 tokens with frequency interpolation
+# Built on DeepSeek MLA + RoPE with Hybrid Muon+AdamW Optimizer
+# Full BF16 precision (no quantization)
+
+# ============================================================================
+# Model Architecture - TaoNet (DeepSeek MLA + RoPE) with YaRN Extension
+# ============================================================================
+model:
+  architecture_type: taonet
+  vocab_size: 8192
+  hidden_dim: 512
+  num_layers: 12
+  num_heads: 8
+  max_seq_length: 8192  # Extended from 1024 → 8192 (8x longer context)
+  
+  # TaoNet-specific: Multi-head Latent Attention (MLA) compression
+  # KV dimension reduced from 512 to 384 (25% memory savings)
+  d_latent_kv: 384
+  
+  # RoPE (Rotary Position Embedding) dimension per head
+  # Default would be 512 / 8 = 64
+  d_rope: 64
+  
+  # Feed-forward intermediate dimension
+  # Default would be 4 * 512 = 2048
+  hidden_dim_ff: 1024
+  
+  # Dropout rate (low for stability)
+  dropout: 0.02
+  
+  # Grouped Query Attention (1 = standard MLA, >1 = GQA)
+  gqa_groups: 1
+  
+  # Optional: Use factorized embedding for parameter efficiency
+  use_factorized_embedding: false
+  d_embed_rank: 96
+  
+  # Weight initialization standard deviation
+  init_std: 0.02
+  
+  # ========================================================================
+  # YaRN (Yet another RoPE eXtension) Configuration
+  # ========================================================================
+  # Enables frequency interpolation to extend context length from 1024 → 8192
+  # The model learns to "pack" RoPE frequencies into the new longer context during training.
+  
+  # RoPE base scale factor (explicit, previously hardcoded to 40)
+  rope_scale: 40.0
+  
+  # Enable YaRN frequency interpolation
+  yarn_enabled: true
+  
+  # Interpolation smoothness parameter
+  # - 1.0 (default): Smooth, gradual interpolation—safer for learning extended context
+  # - 0.5: Aggressive interpolation—faster context expansion, higher risk
+  # - 2.0: Conservative interpolation—safer but slower adaptation
+  # Recommendation: Start with 1.0; tune in follow-up runs if convergence issues
+  yarn_alpha: 1.0
+
+# ============================================================================
+# Dataset Configuration - Local JSONL (Same as Pretrain)
+# ============================================================================
+dataset:
+  local: true
+  jsonl_path: /home/student/Data/TaoData/output.jsonl
+  text_field: text
+  max_samples: 50000  # Reduced from 1M → 50k for quick YaRN adaptation
+  samples_per_chunk: 1000
+  
+  # Tokenizer configuration (unchanged)
+  tokenizer_type: sentencepiece
+  tokenizer_path: tokenizer/tokenizer.model
+  tokenizer_threads: 4
+
+# ============================================================================
+# Training Hyperparameters - Conservative for Context Extension
+# ============================================================================
+# Strategy: Lower learning rates + smaller batch to prevent catastrophic forgetting
+# while the model learns to use 8x longer context.
+
+batch_size: 16  # Reduced from 32 (8192 tokens/seq is memory-intensive)
+num_epochs: 1   # 50k samples / effective_batch=256 ≈ 200 updates (1 epoch sufficient for warm-start)
+
+# Gradient accumulation to maintain effective batch size of ~256
+# Effective batch = batch_size × gradient_accumulation_steps = 16 × 16 = 256
+gradient_accumulation_steps: 16
+
+# Maximum gradient norm for clipping
+max_grad_norm: 1.0
+
+# ============================================================================
+# Optimizer - Hybrid Muon + AdamW (Conservative LR for Stability)
+# ============================================================================
+# Strategy: Use 1/2 of pretrain learning rates to:
+#   1. Avoid catastrophic forgetting of learned features
+#   2. Allow smooth adaptation to YaRN-scaled RoPE frequencies
+#   3. Give the model time to learn how to use extended context
+
+optimizer:
+  optimizer_type: hybrid_muon_adamw
+  
+  # Learning rate for Muon (2D weight matrices)
+  # Reduced: 5e-3 → 2.5e-3 (50% of pretrain)
+  learning_rate: 2.5e-3
+  
+  # Learning rate for AdamW (1D parameters)
+  # Reduced: 5e-4 → 1.25e-4 (25% of pretrain, maintains 1/10 ratio)
+  adamw_lr: 1.25e-4
+  
+  # L2 regularization (weight decay)
+  weight_decay: 0.01
+  
+  # Adam betas (unchanged)
+  betas: [0.9, 0.999]
+  
+  # Epsilon for numerical stability
+  eps: 1e-8
+
+# ============================================================================
+# Learning Rate Scheduler - 3-Phase Cosine with Warmup (Same as Pretrain)
+# ============================================================================
+# Phases:
+# 1. Warmup: 0 → 1.0 (300 steps, ~1.4% of training)
+# 2. Steady: 1.0 (constant for 5% of training steps at peak LR)
+# 3. Decay: 1.0 → 0.1 (cosine decay for remaining ~95%)
+
+scheduler:
+  scheduler_type: cosineWarmup
+  warmup_steps: 300
+  warmup_ratio: 0.0  # Ignored if warmup_steps > 0
+  steady_ratio: 0.05  # 5% of total training steps at peak LR
+  min_lr_ratio: 0.1   # Decay to 10% of peak LR
+  num_cycles: 0.5  # For compatibility (not used in 3-phase schedule)
+
+# ============================================================================
+# Data Type and Device
+# ============================================================================
+dtype: bfloat16  # Use BF16 for better convergence with extended context
+device: cuda     # Use GPU for training
+
+# ============================================================================
+# Checkpointing and Validation
+# ============================================================================
+# Load pretrained checkpoint and continue training
+checkpoint_path: checkpoints/pretrain_taonet/best_model.pt
+checkpoint_dir: checkpoints/yarn_taonet
+save_every_steps: 512    # More frequent saves for 50k samples (200 updates total)
+save_best_model: true
+keep_last_n_checkpoints: 3
+
+# Validation every 512 steps (10% of 50k samples)
+eval_every_steps: 512
+eval_samples: 2500  # Reduced from 8000
+
+# ============================================================================
+# Logging
+# ============================================================================
+log_every_steps: 50  # Log every 50 updates
+aim_repo: .aim
+
+# ============================================================================
+# Miscellaneous
+# ============================================================================
+seed: 42
+num_workers: 0
+pin_memory: true
+
+# ============================================================================
+# YaRN Performance Notes
+# ============================================================================
+# Expected memory usage: ~1.5x of pretrain (8x longer seq, half batch)
+# Expected training time: ~50-100 steps/min on H100 (depends on setup)
+# Expected convergence: Loss should decrease over 50k samples; monitor perplexity on 8192-length sequences
+# 
+# Tuning recommendations for iterative improvements:
+# 1. If loss is unstable: Reduce learning_rate further (1.25e-3)
+# 2. If loss plateaus quickly: Increase max_samples (100k-200k)
+# 3. If memory OOM: Reduce batch_size to 8 (maintain grad_accum at 16)
+# 4. To speed context expansion: Reduce yarn_alpha to 0.5 (more aggressive)
+# 5. For safer training: Increase yarn_alpha to 2.0 (more conservative)

tokenization_taonet.py

@@ -2,6 +2,7 @@
 
 import json
 import os
+import re
 import shutil
 
 from transformers import PreTrainedTokenizer
@@ -43,6 +44,7 @@ class TaoNetTokenizer(PreTrainedTokenizer):
                 str(token): int(token_id) for token, token_id in metadata.get("special_tokens", {}).items()
             }
             configured_special_tokens = [str(token) for token in metadata.get("configured_special_tokens", [])]
+        self.configured_special_tokens = list(configured_special_tokens)
         self.id_to_special_token = {
             int(token_id): str(token) for token, token_id in self.special_token_ids.items()
         }
@@ -74,6 +76,35 @@ class TaoNetTokenizer(PreTrainedTokenizer):
     def _tokenize(self, text):
         return list(self.sp_model.encode(text, out_type=str))
 
+    def get_special_token_id(self, token):
+        return self.special_token_ids.get(token)
+
+    def _encode_with_registered_special_tokens(self, text):
+        if not text:
+            return []
+
+        special_tokens = [
+            token
+            for token in self.configured_special_tokens
+            if token and token in text
+        ]
+        if not special_tokens:
+            return list(self.sp_model.encode(text, out_type=int))
+
+        pattern = "(" + "|".join(re.escape(token) for token in sorted(special_tokens, key=len, reverse=True)) + ")"
+        parts = re.split(pattern, text)
+
+        encoded = []
+        for part in parts:
+            if not part:
+                continue
+            special_token_id = self.special_token_ids.get(part)
+            if special_token_id is not None:
+                encoded.append(int(special_token_id))
+            else:
+                encoded.extend(self.sp_model.encode(part, out_type=int))
+        return encoded
+
     def _convert_token_to_id(self, token):
         if token in self.special_token_ids:
             return self.special_token_ids[token]
@@ -99,6 +130,86 @@ class TaoNetTokenizer(PreTrainedTokenizer):
             return list(token_ids_0)
         return list(token_ids_0) + list(token_ids_1)
 
+    def encode(self, text, return_tensors=None, **kwargs):
+        import torch
+
+        add_special_tokens = kwargs.pop("add_special_tokens", True)
+        del add_special_tokens
+
+        input_ids = self._encode_with_registered_special_tokens(text)
+        if return_tensors == "pt":
+            return torch.tensor(input_ids, dtype=torch.long).unsqueeze(0)
+        return input_ids
+
+    def __call__(self, text, return_tensors=None, **kwargs):
+        import torch
+
+        add_special_tokens = kwargs.pop("add_special_tokens", True)
+        del add_special_tokens
+
+        is_single = isinstance(text, str)
+        texts = [text] if is_single else list(text)
+        encoded_batch = [self._encode_with_registered_special_tokens(item) for item in texts]
+
+        padding = kwargs.pop("padding", False)
+        truncation = kwargs.pop("truncation", False)
+        max_length = kwargs.pop("max_length", None)
+        return_attention_mask = kwargs.pop("return_attention_mask", True)
+
+        if truncation and max_length is not None:
+            encoded_batch = [ids[:max_length] for ids in encoded_batch]
+
+        if padding == "max_length" and max_length is None:
+            raise ValueError("max_length must be set when padding='max_length'")
+        if padding:
+            target_length = max_length if max_length is not None else max(len(ids) for ids in encoded_batch)
+            padded_batch = []
+            attention_masks = []
+            for ids in encoded_batch:
+                trimmed = ids[:target_length]
+                pad_len = target_length - len(trimmed)
+                padded_batch.append(trimmed + [self.pad_token_id] * pad_len)
+                attention_masks.append([1] * len(trimmed) + [0] * pad_len)
+        else:
+            padded_batch = encoded_batch
+            attention_masks = [[1] * len(ids) for ids in encoded_batch]
+
+        if return_tensors == "pt":
+            if not padding and len({len(ids) for ids in padded_batch}) > 1:
+                raise ValueError("Batch elements must have the same length when return_tensors='pt' without padding")
+            input_ids = torch.tensor(padded_batch, dtype=torch.long)
+            result = {"input_ids": input_ids}
+            if return_attention_mask:
+                result["attention_mask"] = torch.tensor(attention_masks, dtype=torch.long)
+            if is_single:
+                result = {key: value for key, value in result.items()}
+            return result
+
+        result = {"input_ids": padded_batch[0] if is_single else padded_batch}
+        if return_attention_mask:
+            result["attention_mask"] = attention_masks[0] if is_single else attention_masks
+        return result
+
+    def build_chat_inputs(self, prompt, return_tensors=None):
+        import torch
+
+        user_token_id = self.special_token_ids["<user>"]
+        assistant_token_id = self.special_token_ids["<assistant>"]
+        prompt_ids = self._encode_with_registered_special_tokens(prompt)
+        input_ids = [user_token_id, *prompt_ids, assistant_token_id]
+        attention_mask = [1] * len(input_ids)
+
+        if return_tensors == "pt":
+            return {
+                "input_ids": torch.tensor(input_ids, dtype=torch.long).unsqueeze(0),
+                "attention_mask": torch.tensor(attention_mask, dtype=torch.long).unsqueeze(0),
+            }
+
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+        }
+
     def save_vocabulary(self, save_directory, filename_prefix=None):
         if not os.path.isdir(save_directory):
             raise ValueError(f"Vocabulary path should be a directory, got: {save_directory}")