rlx_pkg/
from_pt.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Import PyTorch `torch.save` checkpoints (`.pt` / `.pth` / `pytorch_model.bin`)
6//! into an `.rlxp` package.
7
8use crate::tier::StorageTier;
9use crate::write::{ContainerKind, PackedWeight, WriteOptions, write_package};
10use anyhow::{Context, Result};
11use rlx_ir::{DType, Graph, Shape};
12use rlx_nemo::PtModel;
13use std::path::Path;
14
15/// Options for [`pt_to_rlxp`].
16#[derive(Debug, Clone)]
17pub struct PtImportOptions {
18    pub container: ContainerKind,
19    pub include_graph: bool,
20    pub compress_sidecars: bool,
21    pub auto_tier: bool,
22}
23
24impl Default for PtImportOptions {
25    fn default() -> Self {
26        Self {
27            container: ContainerKind::Flat,
28            include_graph: false,
29            compress_sidecars: true,
30            auto_tier: true,
31        }
32    }
33}
34
35/// Convert a PyTorch checkpoint → `.rlxp` (dense f32 weights).
36pub fn pt_to_rlxp(
37    pt_path: impl AsRef<Path>,
38    out: impl AsRef<Path>,
39    opts: &PtImportOptions,
40) -> Result<()> {
41    let pt_path = pt_path.as_ref();
42    let m = PtModel::open(pt_path).with_context(|| format!("open pt {}", pt_path.display()))?;
43
44    let mut weights = Vec::new();
45    for name in m.names() {
46        let t = m.tensor(&name)?;
47        let mut bytes = Vec::with_capacity(t.data.len() * 4);
48        for x in &t.data {
49            bytes.extend_from_slice(&x.to_le_bytes());
50        }
51        weights.push(PackedWeight {
52            name,
53            shape: t.shape,
54            scheme: "f32".into(),
55            layout: "row_major".into(),
56            data: bytes,
57            rank: None,
58            tier: StorageTier::Hot,
59        });
60    }
61
62    let wopts = WriteOptions {
63        name: pt_path
64            .file_stem()
65            .and_then(|s| s.to_str())
66            .unwrap_or("pt")
67            .into(),
68        producer: Some("rlx-pkg/pt_import".into()),
69        container: opts.container,
70        compress_sidecars: opts.compress_sidecars,
71        include_graph: opts.include_graph,
72        write_checksums: true,
73        ..WriteOptions::default()
74    };
75
76    if opts.auto_tier {
77        crate::auto_tier::apply_auto_tier(&mut weights, &Default::default());
78    }
79
80    let graph = if opts.include_graph {
81        let mut g = Graph::new(&wopts.name);
82        let s = Shape::new(&[1], DType::F32);
83        let x = g.input("x", s);
84        g.set_outputs(vec![x]);
85        g
86    } else {
87        Graph::new(&wopts.name)
88    };
89
90    write_package(out.as_ref(), &graph, &weights, &wopts)
91        .with_context(|| format!("write {}", out.as_ref().display()))?;
92    Ok(())
93}
rlx_pkg/from_pt.rs

rlx_pkg/
from_pt.rs
