tracel-ai

PyTorchモデル移植でハマった「自動リナンバー」の罠 TL;DR burn-store の PytorchStore には、パスの数値インデックスを 0, 1, 2, ... に詰め直す map_indices_contiguous という処理があり、デフォルトで有効 PyTorch 側のキーが 0, 2 のように飛び飛びで、Burn 側もその飛びを再現している場合、この詰め直しで重みが別の層にずれる しかもこの処理は with_key_remapping の後に適用されるため、手動リマップでは回避できない 対処は1行です let store = PytorchStore::from_file(&weight_path) .map_indices_contiguous(false); 環境 burn : 0.21.0 burn-store : 0.21.0 発生した症状 PyTorch のモデルを Rust の Burn フレームワークへ移植する作業の中で、PyTorch 側で保存した重み（.pth ファイル）を burn-store 経由で Burn のモデルにロードしようとしたところ、奇妙な失敗の仕方をしました。 キーの不一致でロードがエラーになる。 with_key_remapping で該当するキーを書き換えても、同じエラーが発生して成功しない。 どうにかロードを通過させても、PyTorch 側の出力と数値が一致しない。 前提: パラメータを持たない層はキーを残さない まず、PyTorch 側の事情です。 class MyModule(torch.nn.Module): def init(self): super().init() self.layers = torch.nn.Sequential( torch.nn.Linear(32, 32), torch.nn.ReLU(), torch.nn.Linear(32, 32), torch.nn.ReLU(), ) def forward(self, x): return self.layers(x) torch.nn.Sequential のインデックスは全層に振られますが、ReLU のようなパラメータを持たない層は state_dict にキーを残しません。そのため、保存されるキーは飛び飛びになります。 layers.0.weight / layers.0.bias # 1つ目の Linear layers.2.weight / layers.2.bias # 2つ目の Linear（1 は ReLU なので不在） Burn 側の設計は2通りある このモジュールを Burn へ移植するとき、層の持ち方は大きく分けて2通りあります。 A. パラメータを持つ層だけを詰めて持つ // 概念図: Linear だけを Vec に入れる struct MyModule<B: Backend> { layers: Vec<Linear<B>>, // layers.0, layers.1 を期待する } B. PyTorch の Sequential の並びをそのまま再現する // 概念図: パラメータを持たない層も列挙型のバリアントとして並べる enum Layer<B: Backend> { Linear(Linear<B>), ReLU(PhantomData<B>), } struct MyModule<B: Backend> { layers: Vec<Layer<B>>, // layers.0, layers.2 を期待する } A であれば PyTorch の 0, 2 は Burn の 0, 1 に対応してほしいですし、B であれば 0, 2 のまま入ってほしい。この「どちらを期待しているか」が、今回の話の分かれ目です。 原因: burn-store が勝手にインデックスを詰めていた burn-store 0.21.0 のソースを確認したところ、PytorchStore には map_indices_contiguous というオプションがあり、デフォルトは true でした。ドキュメントコメントには、意図もはっきり書かれています。 When enabled, non-contiguous numeric indices in tensor paths are renumbered to be contiguous. This is useful when loading PyTorch models that have gaps in layer numbering, such as when using nn.Sequential with mixed layer types (e.g., Conv2d layers at indices 0, 2, 4 with ReLU layers at 1, 3, 5). つまり、先ほどの設計Aのための親切機能です。親（プレフィックス）ごとに数値インデックスを集め、ソートして 0 から振り直します。 最小例に適用すると、こうなります。 変換前（PyTorch の生キー） 変換後（Burn に渡るキー） layers.0.* → layers.0.* layers.2.* → layers.1.* # 2つ目の Linear が 1 番に詰められる インデックスがより大きく飛んでいる場合も同じで、たとえば {0, 2, 5, 7} は {0, 1, 2, 3} に変換されます。 設計Bのように Burn 側が飛びを再現していると、Burn が期待する layers.2 は存在せず、代わりに意味の違うキーが並びます。これが、ロードエラーとキー不一致の正体でした。 with_key_remap は解決策にならない .pthのキーと burnモジュールのキーが違うならと、 with_key_remap で解決する方法が思い浮かぶかもしれません。しかし、 with_key_remap ではこの問題を回避できません。 PytorchStore がスナップショットのキャッシュを構築する処理（ensure_snapshots_cache）は、次の順番になっています。 ユーザー定義の with_key_remapping（手動リマップ）が適用される その後に map_indices_contiguous（自動詰め）が適用される つまり、手動で .1 -> .2 のようにどれだけ丁寧に書き換えても、直後に詰め直しが入って、せっかく揃えたインデックスが元通りになってしまいます。 解決策 自動詰めを明示的にオフにします。 // map_indices_contiguous(false) を指定して自動詰めを無効化 let weight_store = PytorchStore::from_file(&weight_path) .map_indices_contiguous(false); これで PyTorch 側の生キーが維持され、 Burn モデルにロードできるようになります。 注意点 常にオフが正解ではない この機能は、設計A（Burn 側がパラメータを持つ層だけを詰めて持つ）では便利です。判断基準はひとつで、「Burn 側のパスが、PyTorch 側の飛びを再現しているか」 です。 Burn 側の構造 設定 飛びを再現している（活性化なども並べる） map_indices_contiguous(false) パラメータを持つ層だけを詰めている デフォルト（true）でよい 設定はストア全体に適応される map_indices_contiguous はモジュール単位ではなく PytorchStore 全体の設定です。1つのストアを複数のモジュールで共有している場合、あるモジュールのためにオフにすると、他のモジュールのキーにも同じく詰め直しが効かなくなります。オフにする場合は、ストアを使う全モジュールが「生キーのまま読む」前提になっているかを確認してください。 まとめ burn-store の map_indices_contiguous はデフォルト有効で、PyTorch の飛び飛びのインデックスを自動で詰める これは Burn 側が層を詰めて持つ設計では便利だが、飛びを再現する設計では重みを黙ってずらす。 with_key_remap は自動詰めより先に適用されるため、この問題の回避策にはならない。 対処は .map_indices_contiguous(false) 。ただしストア全体に適応されるので、共有している他のモジュールも生キー前提か確認が要る。 参考 github.com/tracel-ai/burn/blob/v0.21.0/crates/bur… docs.rs

Beyond Bigger Models: Building a Rust-Based AI Platform with Specialized Nano-LLMs --- description: "Exploring Rust-native AI, specialized nano-models, LoRA adapters, knowledge distillation, RAG, and efficient inference." tags: ai, rust, machinelearning, llm --- Why the future of efficient AI might be an intelligent network of specialized models rather than one enormous neural network. Introduction: Do We Really Need Billions of Parameters for Every Task? The AI industry has spent years pursuing increasingly powerful language models. More parameters. Larger context windows. More compute. More expensive infrastructure. But what if the next major efficiency improvement doesn't come from building another massive model? What if it comes from using smaller, highly specialized models that are exceptionally good at a narrow range of tasks? Consider a public administration system. A citizen submits a question about the deadline for appealing an administrative decision. Does this task really require a general-purpose model with nine billion parameters? Not necessarily. The system needs to: 1. Recognize the administrative procedure. 2. Extract relevant dates and information. 3. Identify the applicable legal framework. 4. Calculate the deadline correctly. 5. Explain the result in understandable language. Most of these operations do not require the full capabilities of a large language model. This observation leads to an architectural concept I want to explore at A:Team Digital Work Solutions: A modular, Rust-based AI platform combining specialized nano-models, shared base models, LoRA adapters, knowledge graphs, and deterministic execution. The objective: Use the smallest amount of computational intelligence necessary to solve a task reliably. --- 1. The Architecture: Specialized Intelligence Instead of One Giant Model The proposed architecture consists of several cooperating components. At the entry point, a semantic router analyzes incoming requests and identifies: - User intent - Domain and subject area - Task complexity - Required knowledge sources - Risk level - Appropriate model or tool The router then selects the most suitable execution path. | Domain | Specialized Component | Responsibilities | |---|---|---| | Administrative Law | Nano-Legal | Appeals, deadlines, administrative documents | | Public Finance | Nano-Budget | Municipal budgets, financial terminology, budget classifications | | General Reasoning | General LLM | Complex analysis, general questions, cross-domain reasoning | | Deterministic Tasks | Rust Tools | Date calculations, financial calculations, validation | The general-purpose model remains available but is not automatically involved in every request. An important architectural decision is that specialized models do not necessarily need independent sets of model weights. Instead, several specialization strategies can coexist. --- 2. Shared Base Models and LoRA Adapters One particularly interesting approach involves using a common base model with multiple specialized LoRA adapters. Low-Rank Adaptation (LoRA) allows us to fine-tune a relatively small number of additional parameters instead of retraining every parameter in the original model. Imagine a shared language model with separate adapters for: - Administrative appeals and procedural law - Municipal budgeting - Taxation and local fees - Public-sector document classification The runtime selects an adapter based on the incoming task. Conceptually: Incoming Request | Semantic Router | Shared Base Model | +----+----+ | | | Legal Budget General LoRA LoRA Mode | | | +----+----+ | Validated Output This architecture could reduce the storage and maintenance overhead of maintaining completely independent fine-tuned models. However, there is an important limitation. LoRA adapters do not automatically reduce the inference cost of the underlying base model. A 9B model with a small adapter still requires the computational resources of approximately a 9B model. That is why the platform should also support genuinely smaller models. --- 3. Nano-LLMs: Small Models for Narrow Problems The second strategy involves specialized models with significantly fewer parameters. Potential experimental model sizes could range from 100 million to one billion parameters, depending on task complexity. Nano-Legal A compact model specialized in: - Identifying administrative procedures - Extracting relevant dates - Recognizing missing information - Classifying administrative documents - Generating structured requests for legal validation - Explaining validated results Nano-Budget A model trained to: - Understand municipal budget terminology - Classify financial information - Identify budget categories - Interpret structured financial datasets - Explain financial terminology - Interact with deterministic financial calculation tools These models would not necessarily replace larger models. Instead, they would handle well-defined tasks independently and escalate difficult cases when required. The actual model size would be determined through benchmarking rather than arbitrary parameter targets. --- 4. Training Specialized Models I see three complementary approaches. Approach A: Fine-Tuning Existing Models The most practical starting point is to specialize an existing small language model. Training data could include: - Structured administrative examples - Domain-specific questions - Document classifications - Validated procedural workflows - Structured input/output examples - Expert-reviewed explanations Depending on model size and available hardware, training could use LoRA or full fine-tuning. Approach B: Knowledge Distillation Another promising approach is knowledge distillation. A larger teacher model generates examples, explanations, and structured task outputs. These are independently validated and used to train a smaller student model. The objective is not to reproduce every capability of the teacher. Instead, the student should learn a narrowly defined set of skills. For example, a legal nano-model might become highly effective at recognizing the information required for an administrative deadline calculation. The actual calculation would still be performed by deterministic software. Approach C: Training from Scratch A longer-term research direction involves developing custom small Transformer architectures. This would require implementing and experimenting with: - Tokenization - Attention mechanisms - Model architecture - Optimization algorithms - Training pipelines - Evaluation datasets - Quantization - Memory management Training from scratch offers greater architectural freedom, but it also introduces substantial data and compute requirements. For that reason, it should follow initial experiments with existing pretrained models. --- 5. Why Rust? Rust is particularly interesting for building the execution platform. It offers: - Native performance - Strong memory-safety guarantees - Efficient concurrency - Low runtime overhead - Predictable resource management - Cross-platform deployment - Integration with existing GPU libraries Existing projects already demonstrate its potential: - Burn — a Rust deep-learning framework. - Candle — a lightweight machine-learning framework. - mistral.rs — a high-performance LLM inference engine. Rather than recreating the entire machine-learning ecosystem, the proposed platform could build on these technologies. A possible modular structure: ateam-ai/ ├── crates/ │ ├── ateam-tensor/ │ ├── ateam-train/ │ ├── ateam-runtime/ │ ├── ateam-router/ │ ├── ateam-adapters/ │ ├── ateam-knowledge/ │ ├── ateam-optimize/ │ └── ateam-eval/ ├── models/ │ ├── nano-legal/ │ ├── nano-budget/ │ └── general/ └── benchmarks/ The first implementation would focus on inference, adapter management, evaluation, and fine-tuning integration. Custom GPU kernels and deeper training-framework development would follow only where measurements demonstrate clear advantages. Rust itself does not automatically make neural networks faster than optimized PyTorch or CUDA implementations. The advantage must come from the complete system architecture and measurable improvements. --- 6. A Practical Example: Administrative Appeal Deadlines Consider a request: > I received an administrative notice dated October 1, 2026. When does my appeal period end? A conventional LLM might attempt to calculate the deadline directly. That creates several risks. The applicable legal framework may differ by jurisdiction, notification method, and procedural context. Instead, the proposed system would use a hybrid architecture. Step 1: Intent Recognition The nano-model identifies the request as an administrative appeal deadline question. Step 2: Information Extraction It extracts: - Document date - Jurisdiction - Notification method - Relevant procedural information - Missing or ambiguous details Step 3: Legal Knowledge Retrieval The system retrieves the relevant legal provisions from a versioned knowledge base. Step 4: Deterministic Calculation A Rust function calculates the deadline using the applicable rules, including calendar adjustments and exceptions. Step 5: Validation The result is checked against required information and legal constraints. Step 6: Explanation The specialized model explains the calculation and cites the applicable legal sources. The critical principle is: Language models should interpret and explain. Deterministic software should perform calculations that require exactness. This separation is especially important in legal, financial, and administrative applications. --- 7. Knowledge Graphs and RAG Specialized training alone is not sufficient. Legal provisions change. Municipal financial data changes. Administrative regulations evolve. Continuously retraining models to reflect every update would be inefficient and potentially unreliable. Therefore, the architecture includes a shared intelligence layer combining: - Retrieval-Augmented Generation (RAG) - Knowledge graphs - Structured databases - Versioned legal documents - Deterministic tools - Output validation Fine-tuning teaches a model how to perform a task. RAG and knowledge graphs provide the current information needed to perform that task. This separation allows specialized models to remain compact while accessing up-to-date knowledge. --- 8. Reducing GPU Time and Inference Costs The primary goal is not merely to reduce parameter counts. It is to reduce the total computational cost of completing useful tasks correctly. Several techniques could contribute. Quantization Quantization reduces model weight precision and memory requirements, potentially with some quality trade-offs. Prefix Caching Prefix caching avoids repeatedly processing identical prompt prefixes. This can be particularly useful in administrative systems where many requests share the same system instructions and procedural context. Continuous Batching Continuous batching improves GPU utilization when multiple requests are processed concurrently. Speculative Decoding Speculative decoding can accelerate generation by using a smaller draft model whose proposed tokens are verified by a larger model. Performance depends on the hardware, model combination, and token acceptance rate. Adaptive Model Routing Simple requests are directed to smaller models. More complex tasks are escalated to larger models. Adapter Sharing Multiple specialized adapters can share the same underlying model. This reduces the need to store and maintain multiple complete fine-tuned models, although it does not inherently reduce base-model computation. The platform should evaluate these techniques individually and in combination. --- 9. How Should We Measure Success? Tokens per second are useful, but they do not tell the complete story. A faster model that produces more incorrect answers is not necessarily more efficient. I would evaluate the platform using the following metrics: | Metric | Purpose | |---|---| | Task accuracy | Measures correctness | | Time to First Token | Measures responsiveness | | Tokens per second | Measures generation throughput | | GPU-seconds per request | Measures compute consumption | | Peak VRAM | Measures memory requirements | | Energy per completed task | Measures energy efficiency | | Cost per validated task | Measures practical operating efficiency | | Escalation rate | Measures dependence on larger models | The central optimization target should be: Minimize computational cost per correctly completed task while maintaining defined quality and reliability thresholds. For legal and administrative applications, incorrect answers must be explicitly measured and cannot be hidden inside an average quality score. --- 10. A Practical Development Roadmap I would approach the project in four stages. Phase 1: Establish the Baseline Run existing small and general-purpose models on a fixed collection of administrative tasks. Measure: - Accuracy - Latency - GPU utilization - Memory consumption - Cost per completed task Phase 2: Build the Rust Runtime Implement: - Semantic routing - Model execution - Adapter management - Deterministic tools - Evaluation and monitoring Reuse established Rust machine-learning libraries wherever practical. Phase 3: Train Specialized Models Develop a nano-model for administrative appeals and another for municipal budgeting. Compare: - Shared-base LoRA specialization - Standalone small models - Knowledge distillation - Different quantization levels Phase 4: Optimize the Complete System Experiment with: - Quantization - Caching - GPU scheduling - Model routing - Memory management - Speculative decoding Measure whether each change actually reduces costs without compromising task quality. --- 11. The Bigger Picture I believe the future of AI infrastructure may involve more than simply deploying increasingly large models. We may see architectures where multiple specialized components collaborate, share knowledge, and execute tasks according to their actual computational requirements. A small model might classify an administrative document. A deterministic Rust function might calculate a deadline. A knowledge graph might identify the relevant legal provisions. A larger language model might only become involved when deeper reasoning is required. The result could be an AI platform that is: - Faster - More resource-efficient - Easier to validate - More economical to operate - More flexible - Better suited to specialized domains This is the direction I want to explore with A:Team Digital Work Solutions. Not because smaller models are automatically better. But because the right model, combined with the right tools and the right knowledge, may be substantially more effective than using the largest available model for everything. The future of efficient AI might not be one enormous brain. It might be an intelligent system of specialized capabilities working together. --- This article describes a proposed research and development architecture. Performance improvements and model-size targets are hypotheses that require experimental validation. What do you think? Could specialized nano-models, shared adapters, and Rust-native inference become a practical alternative to monolithic AI deployments?

Why AI Infrastructure Is Moving to Rust ? GitHub just rewrote its Copilot agent runtime in Rust - and you are probably running Rust inside your Python AI code already. So is Rust replacing Python in AI? No: Python stays the language you drive AI with, while Rust replaces the C and C++ engine underneath and the servers around it. In this video: why Python only gives the orders, the memory bug C++ lets through and Rust refuses to compile, four reasons Rust fits AI systems (no garbage collector, real parallel threads, memory layout control, PyO3), real Rust projects in the AI stack, the honest downsides, and who should actually learn Rust. The C++ and Rust snippets are our own examples; the PyO3 snippet is the official example from the PyO3 README (Apache-2.0 / MIT). Copilot performance figures are GitHub's own measurements, excluding model and network time. Chapters: 0:00 Intro 0:13 GitHub's Rust rewrite 0:48 Rust you already run 1:20 Who really does the math 1:42 Two different jobs 2:00 The problem with C++ 2:42 What Rust does instead 3:07 Why Rust fits AI systems 4:13 Rust inside Python 4:35 Rust around AI 4:59 The catch 5:23 Should you learn Rust? Stock footage: Video by Atakan Ozkan on Pexels; Video by Nothing Ahead on Pexels. Sources (checked 2026-10-17): github.blog huggingface/tokenizers safetensors/safetensors openai/tiktoken github.com/pydantic/pydantic/tree/main/pydantic-c… pydantic.dev pola-rs/polars astral-sh/uv astral-sh/ruff qdrant/qdrant lancedb/lancedb huggingface/candle tracel-ai/burn PyO3/pyo3 PyO3/maturin chromium.org microsoft.com docs.python.org doc.rust-lang.org doc.rust-lang.org docs.pytorch.org huggingface.co docs.lancedb.com docs.python.org raw.githubusercontent.com Independent explainers of AI creative tools. Not affiliated with, endorsed by, or sponsored by any company or tool shown. Prices and features change fast - always check the official sites. Narration uses an AI-generated voice. Script and facts are written and checked by the creator. #rust #python #ai #AItools #AIart #AIvideo

tracel-ai / burn: Burn is a next generation tensor library and Deep Learning Framework that doesn't compromise on flexibility, efficiency and portability. ★16043 tracel-ai/burn

that's the instruction alone, with loading out of the way. the full matmul reading from memory is coming up, at 1.6-1.9x our best f16 kernel so far: cubek#789