pytorch model development

This work addresses the lack of transparent, scalable, and deeply PyTorch-integrated open-source tools for post-training large language models, which hinders research iteration and deployment efficiency. We propose a native PyTorch-based, modular post-training framework centered on the principle of “hackability,” offering composable model builders, training recipes, and a distributed training stack that support diverse fine-tuning strategies and hardware configurations. While maintaining high performance and memory efficiency, the framework significantly enhances code transparency and research flexibility. Empirical evaluations demonstrate that it matches or even surpasses mainstream tools such as Axolotl and Unsloth across multiple post-training scenarios, thereby facilitating efficient and reproducible scientific exploration.

Industrial-scale recommendation and ranking models feature highly complex and continuously evolving architectures, rendering traditional optimization approaches—based on manual intervention or module-level rules—difficult to scale. This work proposes the first extensible and customizable operator-level automatic transformation framework integrated into PyTorch 2.x. By leveraging FX intermediate representation, the PT2 compiler, predefined pattern matching, and a greedy search algorithm, the framework achieves general-purpose model optimizations while strictly preserving computational semantics. Evaluated on real-world industrial recommendation models, the approach delivers up to 63% inference speedup, a 6% reduction in peak memory usage, and over 400 seconds of compilation time savings. The implementation has been open-sourced as part of PyTorch 2.x.

Neural network migration across mainstream frameworks (e.g., PyTorch and TensorFlow) remains challenging due to manual reconstruction requirements, poor compatibility, and semantic discrepancies. To address this, we propose a fully automated cross-framework migration method based on a hub-style intermediate representation (Hub IR). Our approach constructs a unified model IR via abstract syntax tree parsing, then performs semantic-aware structural mapping and framework-specific code generation to achieve bidirectional, functionally equivalent model translation. We systematically resolve two core challenges: cross-framework semantic divergence and topological structure mismatch—addressed for the first time in a unified framework. Experimental evaluation on five representative neural networks demonstrates functional equivalence of generated code, over 90% reduction in manual intervention, and substantial improvements in migration reliability and development efficiency.

Existing GFlowNet research lacks a unified, scalable PyTorch implementation framework, hindering the development of novel training objectives, integration with diverse environments, and reproducible benchmarking. To address this, we introduce the first modular, object-oriented open-source GFlowNet library built on PyTorch. Our method standardizes environment interfaces and sampler abstractions, enables plug-and-play loss functions—including trajectory balance (TB), detailed balance (DB), and unnormalized balance (UB)—and decouples state-space representation, action policies, and flow parameterizations to facilitate customization and composability. The framework successfully reproduces multiple state-of-the-art results across canonical benchmarks, substantially lowering the barrier for algorithm validation and extension. The codebase is publicly released and has been widely adopted by the research community.

Existing LLM pretraining frameworks suffer from fragmentation, poor interoperability, and high maintenance overhead, severely hindering systematic evaluation and production deployment of training methodologies. This paper introduces an open-source, PyTorch-native distributed training system designed for models ranging from 10B to 400B parameters. It proposes a novel modular 3D parallelism architecture—integrating data, tensor, and pipeline parallelism—and tightly couples Float8 quantization with SymmetricMemory hardware-software co-design. The system further incorporates elastic scaling, unified checkpointing, and a reproducible experiment platform. Evaluated on the Llama 3.1 series, it achieves 65.08% speedup on 128 GPUs (8B model), an additional 12.59% gain on 256 GPUs (70B), and a further 30% improvement on 512 GPUs (405B), significantly outperforming baseline systems. The framework delivers high performance, strong scalability, and production readiness.

This study addresses the lack of automated Processing-in-Memory (PIM) offloading support in deep learning frameworks and the associated data movement bottlenecks by proposing a compiler-based automated host-PIM placement strategy. Leveraging MLIR and PyTorch compiler technologies, this method overcomes the limitations of fixed operator lists by dynamically profiling the computational intensity and memory access characteristics of loop nests generated through progressive loop order lowering. This enables profile-guided optimization (PGO)-driven automatic offloading decisions. Evaluated across diverse PIM configurations, the proposed approach achieves speedups of 5.1× and 3.6× over CPU-only execution for GPT-J-6B and LLaMA-7B, respectively.

This work addresses the low success rate in automatic migration from PyTorch to JAX, primarily caused by API discrepancies and dynamic execution behaviors. The authors propose an autonomous agent system that integrates in-context learning (ICL) with an execution oracle. By executing the original PyTorch module to capture ground-truth tensor states as immutable references, the system automatically generates test cases to iteratively guide a large language model in refining the translated JAX code. This approach presents the first closed-loop integration of an execution oracle with ICL-driven self-debugging, achieving high numerical equivalence and reliability across frameworks while maintaining low computational overhead. Evaluated on prominent models including SAM, T5, and CodeWhisperer, the method attains a 91% numerical equivalence rate—substantially outperforming both baseline approaches (9%) and instruction-based self-debugging strategies (27%).

This work proposes ExecuTorch, the first end-to-end deployment framework natively integrated with the PyTorch ecosystem, addressing the fragmentation commonly encountered in edge AI deployment. By introducing an extensible backend abstraction, quantization-aware optimizations, and a unified model serialization format, ExecuTorch preserves the original model semantics while seamlessly targeting heterogeneous hardware—from microcontrollers to specialized accelerators—without sacrificing low latency or offline execution capabilities. The framework bridges the gap between research and production workflows, enabling consistent development and efficient deployment across a broad spectrum of devices, ranging from wearables to compute clusters, thereby significantly enhancing both deployment efficiency and cross-platform consistency.

This work addresses the bottleneck in large-scale deep learning training, which often stems from system implementation rather than model architecture. To this end, the authors develop a lightweight, general-purpose training framework natively implemented in C++ and CUDA from first principles, integrating tensor computation, reverse-mode automatic differentiation, an efficient caching allocator, multi-modal distributed execution, and an MLIR-based compiler. This design achieves fine-grained hardware control while preserving modeling simplicity. On an 8-GPU RTX 6000 Ada system, the framework trains a 124-million-parameter GPT-2 model at 407K tokens/s—surpassing PyTorch’s 395K tokens/s—while reducing memory consumption by 22% and achieving a lower validation loss. This represents the first fully native, end-to-end tunable high-performance training system built solely on standard C++ and core CUDA primitives.

This work addresses the inefficiency and high cost of migrating deep learning models across frameworks—such as from TensorFlow to JAX—in large-scale AI systems. To tackle this challenge, the authors propose an automated, multi-agent collaborative migration approach that integrates static code analysis with an AI-driven planner to generate precise migration instructions. A coordinator and encoder work in tandem, leveraging AI-generated, example-driven migration guides to achieve high-fidelity translation without requiring test code. Innovatively, an AI-based evaluator assesses migration quality, establishing a self-reinforcing development loop. Evaluated in real-world, large-scale production environments, the method accelerates framework migration by 6.4–8×, substantially expediting model infrastructure evolution.