diff --git a/scripts/install.sh b/scripts/install.sh index 229a7db..59d72a4 100755 --- a/scripts/install.sh +++ b/scripts/install.sh @@ -5,7 +5,10 @@ #SIA_INSTALL_NO_NOTEBOOK=1 #SIA_INSTALL_NO_CORE=1 #SIA_INSTALL_NO_ITB=1 -#SIA_INSTALL_NO_MISTRAL_INFER=1 +#SIA_INSTALL_NO_MISTRAL_API_INFER=1 +#SIA_INSTALL_NO_MISTRAL_API_TRAIN=1 +#SIA_INSTALL_NO_MISTRAL_LOCAL_INFER=1 +#SIA_INSTALL_NO_MISTRAL_LOCAL_TRAIN=1 #SIA_INSTALL_NO_GEMMA_INFER=1 #SIA_INSTALL_NO_GEMMA_TRAIN=1 @@ -46,10 +49,28 @@ if [ -z "${SIA_INSTALL_NO_ITB}" ]; then /root/venvs/itb/bin/pip install -e /root/sia/tools/itb fi -if [ -z "${SIA_INSTALL_NO_MISTRAL_INFER}" ]; then - echo "Installing venv for mistral inference" - python3 -m venv /root/venvs/mistral_infer - /root/venvs/mistral_infer/bin/pip install -e /root/sia/tools/mistral_infer +if [ -z "${SIA_INSTALL_NO_MISTRAL_API_INFER}" ]; then + echo "Installing venv for mistral API inference" + python3 -m venv /root/venvs/mistral_api_infer + /root/venvs/mistral_api_infer/bin/pip install -e /root/sia/tools/mistral_api_infer +fi + +if [ -z "${SIA_INSTALL_NO_MISTRAL_API_TRAIN}" ]; then + echo "Installing venv for mistral API training" + python3 -m venv /root/venvs/mistral_api_train + /root/venvs/mistral_api_train/bin/pip install -e /root/sia/tools/mistral_api_train +fi + +if [ -z "${SIA_INSTALL_NO_MISTRAL_LOCAL_INFER}" ]; then + echo "Installing venv for mistral local inference" + python3 -m venv /root/venvs/mistral_local_infer + /root/venvs/mistral_local_infer/bin/pip install -e /root/sia/tools/mistral_local_infer +fi + +if [ -z "${SIA_INSTALL_NO_MISTRAL_LOCAL_TRAIN}" ]; then + echo "Installing venv for mistral local training" + python3 -m venv /root/venvs/mistral_local_train + /root/venvs/mistral_local_train/bin/pip install -e /root/sia/tools/mistral_local_train fi if [ -z "${SIA_INSTALL_NO_GEMMA_INFER}" ]; then diff --git a/tools/mistral_infer/pyproject.toml b/tools/mistral_api_infer/pyproject.toml similarity index 80% rename from tools/mistral_infer/pyproject.toml rename to tools/mistral_api_infer/pyproject.toml index 18b1c62..76fe605 100644 --- a/tools/mistral_infer/pyproject.toml +++ b/tools/mistral_api_infer/pyproject.toml @@ -3,7 +3,7 @@ requires = ["setuptools>=42", "wheel"] build-backend = "setuptools.build_meta" [project] -name = "mistral_infer" +name = "mistral_api_infer" version = "0.1.0" requires-python = ">=3.8" @@ -15,4 +15,4 @@ dependencies = [ ] [project.scripts] -mistral_infer = "mistral_infer.__main__:main" \ No newline at end of file +mistral_api_infer = "mistral_api_infer.__main__:main" \ No newline at end of file diff --git a/tools/mistral_infer/src/mistral_infer/__main__.py b/tools/mistral_api_infer/src/mistral_api_infer/__main__.py similarity index 95% rename from tools/mistral_infer/src/mistral_infer/__main__.py rename to tools/mistral_api_infer/src/mistral_api_infer/__main__.py index 1b02275..80e4654 100644 --- a/tools/mistral_infer/src/mistral_infer/__main__.py +++ b/tools/mistral_api_infer/src/mistral_api_infer/__main__.py @@ -1,4 +1,4 @@ -from .mistral_llm_engine import MistralLlmEngine +from .mistral_llm_engine import MistralApiLlmEngine from dotenv import load_dotenv from llm_engine_utils.protocol import process import argparse diff --git a/tools/mistral_infer/src/mistral_infer/mistral_llm_engine.py b/tools/mistral_api_infer/src/mistral_api_infer/mistral_llm_engine.py similarity index 100% rename from tools/mistral_infer/src/mistral_infer/mistral_llm_engine.py rename to tools/mistral_api_infer/src/mistral_api_infer/mistral_llm_engine.py diff --git a/tools/mistral_train/pyproject.toml b/tools/mistral_api_train/pyproject.toml similarity index 82% rename from tools/mistral_train/pyproject.toml rename to tools/mistral_api_train/pyproject.toml index 26d1d58..305ae03 100644 --- a/tools/mistral_train/pyproject.toml +++ b/tools/mistral_api_train/pyproject.toml @@ -3,7 +3,7 @@ requires = ["setuptools>=42", "wheel"] build-backend = "setuptools.build_meta" [project] -name = "mistral_train" +name = "mistral_api_train" version = "0.1.0" requires-python = ">=3.8" @@ -16,4 +16,4 @@ dependencies = [ ] [project.scripts] -mistral_train = "mistral_train.__main__:main" \ No newline at end of file +mistral_api_train = "mistral_api_train.__main__:main" \ No newline at end of file diff --git a/tools/mistral_train/src/mistral_train/__main__.py b/tools/mistral_api_train/src/mistral_api_train/__main__.py similarity index 100% rename from tools/mistral_train/src/mistral_train/__main__.py rename to tools/mistral_api_train/src/mistral_api_train/__main__.py diff --git a/tools/mistral_train/src/mistral_train/config.py b/tools/mistral_api_train/src/mistral_api_train/config.py similarity index 100% rename from tools/mistral_train/src/mistral_train/config.py rename to tools/mistral_api_train/src/mistral_api_train/config.py diff --git a/tools/mistral_local_infer/pyproject.toml b/tools/mistral_local_infer/pyproject.toml new file mode 100644 index 0000000..a0e13b4 --- /dev/null +++ b/tools/mistral_local_infer/pyproject.toml @@ -0,0 +1,25 @@ +[build-system] +requires = ["setuptools>=42", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "mistral_local_infer" +version = "0.1.0" +requires-python = ">=3.8" + +dependencies = [ + "blobfile>=3.0.0", + "llama-cpp-python @ git+https://github.com/abetlen/llama-cpp-python.git@v0.3.16#egg=llama-cpp-python&env=CMAKE_ARGS=-DLLAMA_BUILD=OFF", + "llm_engine_utils @ file:///root/sia/lib/llm_engine_utils", + "mistral-common>=1.8.6", + "protobuf>=6.0.0", + "python-dotenv>=1.0.0", + "sentencepiece>=0.2.0", + "tiktoken>=0.9.0", + "transformers>=5.0.0rc0", + "vulkan", + "xml_schema_validator @ file:///root/sia/lib/xml_schema_validator", +] + +[project.scripts] +mistral_local_infer = "mistral_local_infer.__main__:main" diff --git a/tools/mistral_local_infer/src/mistral_local_infer/__main__.py b/tools/mistral_local_infer/src/mistral_local_infer/__main__.py new file mode 100644 index 0000000..c5b2e27 --- /dev/null +++ b/tools/mistral_local_infer/src/mistral_local_infer/__main__.py @@ -0,0 +1,42 @@ +from .mistral_llm_engine import MistralLlmEngine +from dotenv import load_dotenv +from llm_engine_utils.protocol import process +import argparse +import os + +def main(): + load_dotenv() + parser = argparse.ArgumentParser(description='Ministral-3 Local Inference using llama.cpp with Vulkan') + + parser.add_argument( + '--model', + type=str, + default=os.getenv('SIA_MISTRAL_MODEL', '/root/models/current/model.gguf'), + help='Model name (default: /root/models/current/model.gguf, env: SIA_MISTRAL_MODEL)' + ) + + parser.add_argument( + '--tokenizer', + type=str, + default=os.getenv('SIA_MISTRAL_TOKENIZER', '/root/models/current/tokenizer'), + help='Model name (default: /root/models/current/tokenizer, env: SIA_MISTRAL_TOKENIZER)' + ) + + parser.add_argument( + '--token-limit', + type=int, + default=os.getenv('SIA_MISTRAL_TOKEN_LIMIT', 10000), + help='Token limit (default: 10000, env: SIA_MISTRAL_TOKEN_LIMIT)' + ) + + args = parser.parse_args() + + mistral_llm_engine = MistralLlmEngine( + model=args.model, + tokenizer=args.tokenizer, + token_limit=args.token_limit, + ) + process(mistral_llm_engine) + +if __name__ == "__main__": + main() diff --git a/tools/mistral_local_infer/src/mistral_local_infer/mistral_llm_engine.py b/tools/mistral_local_infer/src/mistral_local_infer/mistral_llm_engine.py new file mode 100644 index 0000000..df07e58 --- /dev/null +++ b/tools/mistral_local_infer/src/mistral_local_infer/mistral_llm_engine.py @@ -0,0 +1,97 @@ +import os +os.environ["LLAMA_CPP_LIB_PATH"] = "/usr/local/lib" +os.environ["LD_LIBRARY_PATH"] += ":/usr/local/lib" +os.chdir("/usr/local/lib") + +from llama_cpp import Llama, LogitsProcessorList, llama_cpp +from llm_engine_utils import LlmEngine +from llm_engine_utils.iterators import skip_prefix +from pathlib import Path +from transformers import AutoTokenizer +from typing import Iterator +from xml_schema_validator import LlamaCppLogitsProcessor + +llama_cpp._lib.ggml_backend_load_all() + +class MistralLlmEngine(LlmEngine): + """ + Ministral-3 inference engine using llama.cpp with Vulkan acceleration. + + Ministral-3 architecture features: + - 34 transformer layers with alternating attention (1 full + 3 sliding window) + - 131K vocabulary tokens + - Supports up to 256K context window + - Uses Grouped Query Attention (32 heads, 8 key-value heads) + """ + def __init__( + self, + model: str, + tokenizer: str, + token_limit: int, + ): + self._model = model + self._tokenizer = AutoTokenizer.from_pretrained(tokenizer) + self._token_limit = token_limit + + self._llm = Llama( + model_path=model, + n_gpu_layers=0, + n_ctx=token_limit, + flash_attn=True, + ) + + def infer_xml(self, schema: Path, system: str, context: str, prefix: str) -> Iterator[str]: + xml_schema_text = Path(schema).read_text() + logits_processor = LlamaCppLogitsProcessor(self._tokenizer, xml_schema_text).get_processor() + logits_processor_list = LogitsProcessorList([logits_processor]) + prompt = self._format_messages(system, context, prefix) + stream = self._llm.create_completion( + prompt=prompt, + max_tokens=self._token_limit, + stream=True, + logits_processor=logits_processor_list + ) + + def content_generator(): + for output in stream: + choice = output["choices"][0] + if choice.get('finish_reason'): + break + if 'text' in choice: + text = choice['text'] + if text == '': + break + yield text + + yield from skip_prefix(content_generator(), prefix) + + def token_count(self, system: str, context: str) -> int: + prompt = self._format_messages(system, context, None) + tokens = self._tokenizer.encode(prompt) + return len(tokens) + + def token_limit(self) -> int: + return self._token_limit + + def _format_messages(self, system: str, context: str, prefix: str) -> str: + messages = [ + { + "role": "user", + "content": f"{system}\n\n--- Context ---\n{context}", + }, + { + "role": "assistant", + "content": prefix, + }, + ] if prefix else [ + { + "role": "user", + "content": f"{system}\n\n--- Context ---\n{context}", + } + ] + return self._tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=not prefix, + continue_final_message=bool(prefix) + ).removeprefix("") diff --git a/tools/mistral_local_train/pyproject.toml b/tools/mistral_local_train/pyproject.toml new file mode 100644 index 0000000..1901c96 --- /dev/null +++ b/tools/mistral_local_train/pyproject.toml @@ -0,0 +1,27 @@ +[build-system] +requires = ["setuptools>=42", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "mistral_local_train" +version = "0.1.0" +requires-python = ">=3.8" + +dependencies = [ + "accelerate==1.2.0", + "bitsandbytes>=0.45.0", + "datasets==3.3.2", + "evaluate==0.4.3", + "kernels>=0.11.1", + "llm_engine_utils @ file:///root/sia/lib/llm_engine_utils", + "mistral-common>=1.8.6", + "peft==0.13.2", + "python-dotenv>=1.0.0", + "sentencepiece>=0.2.0", + "torch==2.4.1", + "transformers>=5.0.0rc0", + "trl>=0.17.0", +] + +[project.scripts] +mistral_local_train = "mistral_local_train.__main__:main" diff --git a/tools/mistral_local_train/src/mistral_local_train/__main__.py b/tools/mistral_local_train/src/mistral_local_train/__main__.py new file mode 100644 index 0000000..6a3ea6b --- /dev/null +++ b/tools/mistral_local_train/src/mistral_local_train/__main__.py @@ -0,0 +1,106 @@ +from llm_engine_utils.dataset import Dataset +from pathlib import Path +from peft import LoraConfig, AutoPeftModelForCausalLM +from transformers import AutoTokenizer, TrainingArguments +from transformers.models.ministral3 import Ministral3ForCausalLM +from trl import SFTTrainer +import os +import sys +import torch + +from .config import Config + +def main(): + config = Config() + train(config) + merge(config) + +def train(config: Config): + tokenizer = AutoTokenizer.from_pretrained( + config.model, + token=config.api_key, + trust_remote_code=True, + use_fast=False, + ) + tokenizer.save_pretrained(config.output_dir/"tokenizer") + + # Use Ministral3ForCausalLM for text-only training + # Note: torch_dtype is deprecated in transformers 5.x, use dtype instead + model = Ministral3ForCausalLM.from_pretrained( + config.model, + dtype=torch.float32, + device_map="cpu", + token=config.api_key, + attn_implementation='eager', + low_cpu_mem_usage=True, + ) + + dataset = Dataset(config.config_path) + dataset.validate() + dataset = dataset.to_transformers_dataset(tokenizer) + + # LoRA config targeting Ministral-3 architecture layers + # Ministral-3 uses alternating attention (1 full + 3 sliding window) across 34 layers + # Target modules include: + # - Attention projections: q_proj, k_proj, v_proj, o_proj + # - MLP projections: gate_proj, up_proj, down_proj + # - Language model head: lm_head + lora_config = LoraConfig( + r=8, + lora_alpha=16, + target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "lm_head"], + lora_dropout=0.05, + bias="none", + task_type="CAUSAL_LM", + ) + + training_args = TrainingArguments( + per_device_train_batch_size=1, + gradient_accumulation_steps=4, + warmup_steps=1, + num_train_epochs=2, + learning_rate=1e-3, + fp16=False, + logging_steps=5, + save_strategy="steps", + save_steps=10, + output_dir=config.output_dir/"lora", + optim="adamw_torch", + seed=42, + group_by_length=True, + use_cpu=True, + ) + + trainer = SFTTrainer( + model=model, + train_dataset=dataset, + args=training_args, + peft_config=lora_config, + formatting_func=format_sia_example, + ) + + trainer.train() + + trainer.model.save_pretrained(config.output_dir/"lora_adapter") + +def merge(config: Config): + # Note: torch_dtype is deprecated in transformers 5.x, use dtype instead + adapted_model = AutoPeftModelForCausalLM.from_pretrained( + config.output_dir/"lora_adapter", + dtype=torch.float16, + token=config.api_key, + load_in_4bit=False, + load_in_8bit=False, + ) + merged_model = adapted_model.merge_and_unload() + + merged_model.save_pretrained( + config.output_dir/"merged", + safe_serialization=True + ) + +def format_sia_example(example): + return example['messages'].removeprefix("") + +if __name__ == "__main__": + exit(main()) diff --git a/tools/mistral_local_train/src/mistral_local_train/config.py b/tools/mistral_local_train/src/mistral_local_train/config.py new file mode 100644 index 0000000..3834b2f --- /dev/null +++ b/tools/mistral_local_train/src/mistral_local_train/config.py @@ -0,0 +1,50 @@ +from pathlib import Path +import argparse +import os + +class Config: + def __init__(self): + parser = argparse.ArgumentParser( + description='Train Ministral-3 model locally using LoRA fine-tuning' + ) + parser.add_argument( + '--config', + type=Path, + default=Path('/root/sia/training/config.yaml'), + help='Path to config file (default: /root/sia/training/config.yaml)' + ) + parser.add_argument( + '--model', + type=str, + default='mistralai/Ministral-3-8B-Base-2512', + help='Base model for fine-tuning (default: mistralai/Ministral-3-8B-Base-2512)' + ) + parser.add_argument( + '--api-key', + type=str, + default=os.environ.get('SIA_HF_API_KEY'), + help='Huggingface API key (optional, env: SIA_HF_API_KEY)' + ) + parser.add_argument( + '--output-dir', + type=Path, + default=Path('/root/models/current'), + help='Output directory for fine-tuned model and converted gguf (default: /root/models/current)' + ) + self.args = parser.parse_args() + + @property + def config_path(self) -> Path: + return self.args.config + + @property + def model(self) -> str: + return self.args.model + + @property + def api_key(self) -> str: + return self.args.api_key + + @property + def output_dir(self) -> Path: + return self.args.output_dir