boostr 0.1.0

ML framework built on numr - attention, quantization, model architectures
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104

//! Core model trait

use crate::error::Result;
use crate::model::config::ModelConfig;
use crate::nn::VarBuilder;
use crate::ops::traits::architecture::moe::MoEOps;
use crate::ops::traits::position::alibi::AlibiOps;
use crate::ops::traits::{FlashAttentionOps, KvCacheOps, PagedAttentionOps, RoPEOps};
use crate::quant::traits::{DequantOps, QuantMatmulOps};
use numr::autograd::Var;
use numr::ops::{
    ActivationOps, BinaryOps, CompareOps, ConditionalOps, IndexingOps, NormalizationOps, ReduceOps,
    ScalarOps, ShapeOps, TensorOps, UnaryOps,
};
use numr::runtime::{Runtime, RuntimeClient};

/// Trait alias for the full set of client bounds required by model forward passes.
pub trait ModelClient<R: Runtime>:
    RuntimeClient<R>
    + TensorOps<R>
    + ScalarOps<R>
    + ReduceOps<R>
    + IndexingOps<R>
    + ShapeOps<R>
    + ActivationOps<R>
    + BinaryOps<R>
    + UnaryOps<R>
    + CompareOps<R>
    + ConditionalOps<R>
    + RoPEOps<R>
    + FlashAttentionOps<R>
    + PagedAttentionOps<R>
    + KvCacheOps<R>
    + QuantMatmulOps<R>
    + NormalizationOps<R>
    + MoEOps<R>
    + AlibiOps<R>
{
}

impl<R, C> ModelClient<R> for C
where
    R: Runtime,
    C: RuntimeClient<R>
        + TensorOps<R>
        + ScalarOps<R>
        + ReduceOps<R>
        + IndexingOps<R>
        + ShapeOps<R>
        + ActivationOps<R>
        + BinaryOps<R>
        + UnaryOps<R>
        + CompareOps<R>
        + ConditionalOps<R>
        + RoPEOps<R>
        + FlashAttentionOps<R>
        + PagedAttentionOps<R>
        + KvCacheOps<R>
        + QuantMatmulOps<R>
        + NormalizationOps<R>
        + MoEOps<R>
        + AlibiOps<R>,
{
}

/// Core trait for all model architectures.
///
/// Models take token IDs and produce logits. All internal computation
/// uses `Var<R>` for autograd compatibility.
pub trait Model<R: Runtime>: Sized {
    /// Create model from configuration with zero-initialized weights.
    fn from_config(config: &ModelConfig, device: &R::Device) -> Result<Self>;

    /// Create model from a VarBuilder (loads real weights).
    ///
    /// Takes `&mut` because tensors are moved out of the VarMap (zero-copy).
    fn from_varbuilder(vb: &mut VarBuilder<R>, config: &ModelConfig) -> Result<Self>
    where
        R::Client: DequantOps<R> + numr::ops::TypeConversionOps<R>,
    {
        let _ = (vb, config);
        Err(crate::error::Error::ModelError {
            reason: "from_varbuilder not implemented for this model".into(),
        })
    }

    /// Forward pass: token_ids [B, S] -> logits [B, S, vocab_size]
    fn forward<C>(&self, client: &C, input_ids: &Var<R>) -> Result<Var<R>>
    where
        C: ModelClient<R>,
        R::Client: TensorOps<R>
            + ScalarOps<R>
            + ReduceOps<R>
            + IndexingOps<R>
            + ShapeOps<R>
            + ActivationOps<R>
            + BinaryOps<R>
            + UnaryOps<R>
            + CompareOps<R>
            + ConditionalOps<R>;

    /// Get the model configuration
    fn config(&self) -> &ModelConfig;
}