burn_core/module/param/
tensor.rs

1use super::{LoraAdapter, Param, ParamId, Parameter};
2use crate::module::{
3    AutodiffModule, Content, Module, ModuleDisplay, ModuleDisplayDefault, ModuleMapper,
4    ModuleVisitor,
5};
6use alloc::{boxed::Box, format, string::ToString, vec::Vec};
7use burn_tensor::{Bool, Device, Float, Int, Tensor, TensorData};
8
9impl<const D: usize> super::sealed::Sealed for Tensor<D, Float> {}
10impl<const D: usize> super::sealed::Sealed for Tensor<D, Int> {}
11impl<const D: usize> super::sealed::Sealed for Tensor<D, Bool> {}
12
13impl<const D: usize> Parameter for Tensor<D, Float> {
14    fn device(&self) -> Device {
15        Tensor::device(self)
16    }
17
18    fn is_require_grad(&self) -> bool {
19        Tensor::is_require_grad(self)
20    }
21
22    fn set_require_grad(self, require_grad: bool) -> Self {
23        Tensor::set_require_grad(self, require_grad)
24    }
25
26    fn shape(&self) -> burn_std::Shape {
27        Tensor::shape(self)
28    }
29
30    fn load_to_device(self, device: &Device) -> Self {
31        if self.device() != *device {
32            Tensor::to_device(self, device).detach()
33        } else {
34            self
35        }
36    }
37
38    fn compose_lora(self, adapter: &LoraAdapter) -> Self {
39        // `delta` has shape `[d_in, d_out]` (rank 2); the base weight is rank `D == 2` when an
40        // adapter is attached, so the reshape is an identity that only adjusts the static rank.
41        let delta = adapter.delta().reshape(Tensor::shape(&self));
42        self.add(delta)
43    }
44}
45
46impl<const D: usize> Parameter for Tensor<D, Int> {
47    fn device(&self) -> Device {
48        Tensor::device(self)
49    }
50
51    fn is_require_grad(&self) -> bool {
52        false
53    }
54
55    fn set_require_grad(self, _require_grad: bool) -> Self {
56        self
57    }
58
59    fn shape(&self) -> burn_std::Shape {
60        Tensor::shape(self)
61    }
62
63    fn load_to_device(self, device: &Device) -> Self {
64        if self.device() != *device {
65            Tensor::to_device(self, device)
66        } else {
67            self
68        }
69    }
70}
71
72impl<const D: usize> Parameter for Tensor<D, Bool> {
73    fn device(&self) -> Device {
74        Tensor::device(self)
75    }
76
77    fn is_require_grad(&self) -> bool {
78        false
79    }
80
81    fn set_require_grad(self, _require_grad: bool) -> Self {
82        self
83    }
84
85    fn shape(&self) -> burn_std::Shape {
86        Tensor::shape(self)
87    }
88
89    fn load_to_device(self, device: &Device) -> Self {
90        if self.device() != *device {
91            Tensor::to_device(self, device)
92        } else {
93            self
94        }
95    }
96}
97
98impl<const D: usize> Param<Tensor<D>> {
99    /// Create a new parameter from a float tensor.
100    ///
101    /// # Warnings
102    ///
103    /// We strongly recommend using [Param::uninitialized] if you are using this method to
104    /// initialize parameters inside a module, since the tensor initialization will be lazy,
105    /// making the loading of weights more performant.
106    pub fn from_tensor(value: Tensor<D>) -> Self {
107        // When creating a parameter from a float tensor, we automatically mark it as requiring
108        // gradients, so that it can be updated by an optimizer.
109        Param::initialized(ParamId::new(), value.require_grad())
110    }
111
112    /// Create a new parameter from data.
113    pub fn from_data<T>(data: T, device: &Device) -> Self
114    where
115        T: Into<TensorData>,
116    {
117        let data: TensorData = data.into();
118        // When creating a parameter from a float tensor, we automatically mark it as requiring
119        // gradients, so that it can be updated by an optimizer.
120        device.memory_persistent_allocations(data, |data| {
121            let value = Tensor::from_data(data, device);
122            Param::initialized(ParamId::new(), value.require_grad())
123        })
124    }
125}
126
127/// Visit the trainable factors of a LoRA [adapter](LoraAdapter) as nested parameters, so the
128/// record/optimizer traversal sees them at stable paths (e.g. `weight.lora.a`, `weight.lora.b`).
129fn visit_adapter<V: ModuleVisitor>(adapter: &LoraAdapter, visitor: &mut V) {
130    visitor.enter_module("lora", "Struct:LoraAdapter");
131    visitor.enter_module("a", "Struct:LoraAdapter");
132    Module::visit(&adapter.a, visitor);
133    visitor.exit_module("a", "Struct:LoraAdapter");
134    visitor.enter_module("b", "Struct:LoraAdapter");
135    Module::visit(&adapter.b, visitor);
136    visitor.exit_module("b", "Struct:LoraAdapter");
137    visitor.exit_module("lora", "Struct:LoraAdapter");
138}
139
140/// Map the trainable factors of a LoRA [adapter](LoraAdapter), mirroring [`visit_adapter`] so the
141/// optimizer/record mapper resolves the same paths.
142fn map_adapter<M: ModuleMapper>(adapter: LoraAdapter, mapper: &mut M) -> LoraAdapter {
143    let LoraAdapter { a, b, scale } = adapter;
144    mapper.enter_module("lora", "Struct:LoraAdapter");
145    mapper.enter_module("a", "Struct:LoraAdapter");
146    let a = Module::map(a, mapper);
147    mapper.exit_module("a", "Struct:LoraAdapter");
148    mapper.enter_module("b", "Struct:LoraAdapter");
149    let b = Module::map(b, mapper);
150    mapper.exit_module("b", "Struct:LoraAdapter");
151    mapper.exit_module("lora", "Struct:LoraAdapter");
152    LoraAdapter { a, b, scale }
153}
154
155impl<const D: usize> Module for Param<Tensor<D>> {
156    fn visit<V: ModuleVisitor>(&self, visitor: &mut V) {
157        match self.adapter() {
158            None => visitor.visit_float(self),
159            Some(adapter) => {
160                // Visit the frozen base, then the trainable adapter factors as separate leaves.
161                visitor.visit_float(&self.without_adapter());
162                visit_adapter(adapter, visitor);
163            }
164        }
165    }
166
167    fn map<M: ModuleMapper>(mut self, mapper: &mut M) -> Self {
168        match self.adapter.take() {
169            None => mapper.map_float(self),
170            Some(adapter) => {
171                // `self` no longer carries the adapter, so the mapper operates on the raw base.
172                let base = mapper.map_float(self);
173                let adapter = map_adapter(*adapter, mapper);
174                base.with_adapter(Some(Box::new(adapter)))
175            }
176        }
177    }
178
179    fn to_device(mut self, device: &Device) -> Self {
180        let adapter = self.adapter.take();
181        let base = self.map(|tensor| tensor.to_device(device));
182        match adapter {
183            None => base,
184            Some(adapter) => base.with_adapter(Some(Box::new(LoraAdapter {
185                a: Module::to_device(adapter.a, device),
186                b: Module::to_device(adapter.b, device),
187                scale: adapter.scale,
188            }))),
189        }
190    }
191
192    fn fork(mut self, device: &Device) -> Self {
193        let adapter = self.adapter.take();
194        let base = self.map(|tensor| {
195            let is_require_grad = tensor.is_require_grad();
196            let mut tensor = tensor.to_device(device).detach();
197
198            if is_require_grad {
199                tensor = tensor.require_grad();
200            }
201
202            tensor
203        });
204        match adapter {
205            None => base,
206            Some(adapter) => base.with_adapter(Some(Box::new(LoraAdapter {
207                a: Module::fork(adapter.a, device),
208                b: Module::fork(adapter.b, device),
209                scale: adapter.scale,
210            }))),
211        }
212    }
213
214    fn collect_devices(&self, mut devices: Vec<Device>) -> Vec<Device> {
215        let device = self.base().device();
216
217        if !devices.contains(&device) {
218            devices.push(device)
219        }
220
221        if let Some(adapter) = self.adapter() {
222            devices = Module::collect_devices(&adapter.a, devices);
223            devices = Module::collect_devices(&adapter.b, devices);
224        }
225
226        devices
227    }
228}
229
230impl<const D: usize> ModuleDisplayDefault for Param<Tensor<D>> {
231    fn content(&self, content: Content) -> Option<Content> {
232        let id = if content.display_settings.show_param_id() {
233            format!(", id: {}", self.id)
234        } else {
235            "".to_string()
236        };
237        let string = format!(
238            "ParamTensor {{rank: {D}, shape: {:?}, kind: float{id}}}",
239            self.shape().as_slice()
240        );
241        content.add_formatted(&string).optional()
242    }
243}
244impl<const D: usize> ModuleDisplay for Param<Tensor<D>> {}
245
246impl<const D: usize> Module for Param<Tensor<D, Int>> {
247    fn visit<V: ModuleVisitor>(&self, visitor: &mut V) {
248        visitor.visit_int(self)
249    }
250
251    fn map<M: ModuleMapper>(self, mapper: &mut M) -> Self {
252        mapper.map_int(self)
253    }
254
255    fn to_device(self, device: &Device) -> Self {
256        self.map(|tensor| tensor.to_device(device))
257    }
258
259    fn fork(self, device: &Device) -> Self {
260        self.to_device(device) // Don't support autodiff.
261    }
262
263    fn collect_devices(&self, mut devices: Vec<Device>) -> Vec<Device> {
264        let device = self.val().device();
265
266        if !devices.contains(&device) {
267            devices.push(device)
268        }
269
270        devices
271    }
272}
273
274impl<const D: usize> ModuleDisplayDefault for Param<Tensor<D, Int>> {
275    fn content(&self, content: Content) -> Option<Content> {
276        let id = if content.display_settings.show_param_id() {
277            format!(", id: {}", self.id)
278        } else {
279            "".to_string()
280        };
281        let string = format!(
282            "ParamTensor {{rank: {D}, shape: {:?}, kind: int{id}}}",
283            self.shape().as_slice()
284        );
285        content.add_formatted(&string).optional()
286    }
287}
288impl<const D: usize> ModuleDisplay for Param<Tensor<D, Int>> {}
289
290impl<const D: usize> Module for Param<Tensor<D, Bool>> {
291    fn visit<V: ModuleVisitor>(&self, visitor: &mut V) {
292        visitor.visit_bool(self)
293    }
294
295    fn map<M: ModuleMapper>(self, mapper: &mut M) -> Self {
296        mapper.map_bool(self)
297    }
298
299    fn to_device(self, device: &Device) -> Self {
300        self.map(|tensor| tensor.to_device(device))
301    }
302
303    fn fork(self, device: &Device) -> Self {
304        self.to_device(device) // Don't support autodiff.
305    }
306
307    fn collect_devices(&self, mut devices: Vec<Device>) -> Vec<Device> {
308        let device = self.val().device();
309
310        if !devices.contains(&device) {
311            devices.push(device)
312        }
313
314        devices
315    }
316}
317
318impl<const D: usize> ModuleDisplayDefault for Param<Tensor<D, Bool>> {
319    fn content(&self, content: Content) -> Option<Content> {
320        let id = if content.display_settings.show_param_id() {
321            format!(", id: {}", self.id)
322        } else {
323            "".to_string()
324        };
325
326        let string = format!(
327            "ParamTensor {{rank: {D}, shape: {:?}, kind: bool{id}}}",
328            self.shape().as_slice()
329        );
330        content.add_formatted(&string).optional()
331    }
332}
333
334impl<const D: usize> ModuleDisplay for Param<Tensor<D, Bool>> {}
335
336impl<const D: usize> AutodiffModule for Param<Tensor<D>> {
337    fn valid(&self) -> Self {
338        // Preserve initialized param `require_grad` state, but reset the inner value's.
339        // When a LoRA adapter is attached, `val()` folds it into the base for inference, so the
340        // resulting inference parameter is a plain (adapter-free) composed weight.
341        let require_grad = self.require_grad;
342        let mut param = Param::initialized(self.id, self.val().inner().set_require_grad(false));
343        param.require_grad = require_grad;
344        param
345    }
346
347    fn from_inner(mut module: Self) -> Self {
348        // Keep the adapter structure (and its trainable factors) when moving onto the autodiff
349        // backend, so the adapter remains trainable after `train()`.
350        let adapter = module.adapter.take();
351        // Reinstate the param's `require_grad` state
352        let tensor = Tensor::from_inner(module.val()).set_require_grad(module.require_grad);
353        let base = Param::initialized(module.id, tensor);
354        match adapter {
355            None => base,
356            Some(adapter) => base.with_adapter(Some(Box::new(LoraAdapter {
357                a: AutodiffModule::from_inner(adapter.a),
358                b: AutodiffModule::from_inner(adapter.b),
359                scale: adapter.scale,
360            }))),
361        }
362    }
363}
364
365// impl<const D: usize, B: AutodiffBackend> HasAutodiffModule
366//     for Param<Tensor<B::InnerBackend, D>>
367// {
368//     type TrainModule = Param<Tensor<D>>;
369// }
370
371impl<const D: usize> AutodiffModule for Param<Tensor<D, Int>> {
372    fn valid(&self) -> Self {
373        Param::initialized(self.id, self.val().inner())
374    }
375
376    fn from_inner(module: Self) -> Self {
377        Param::initialized(module.id, Tensor::from_inner(module.val()))
378    }
379}
380
381impl<const D: usize> AutodiffModule for Param<Tensor<D, Bool>> {
382    fn valid(&self) -> Self {
383        Param::initialized(self.id, self.val().inner())
384    }
385
386    fn from_inner(module: Self) -> Self {
387        Param::initialized(module.id, Tensor::from_inner(module.val()))
388    }
389}
390
391#[cfg(all(test, feature = "std", feature = "autodiff"))]
392mod tests {
393    use super::*;
394    use crate::{module::Module, test_device};
395
396    #[test]
397    fn test_param_require_grad_stateful() {
398        let device = test_device().autodiff();
399        let tensor = Tensor::<2>::ones([3, 3], &device).require_grad();
400
401        let param = Param::initialized(ParamId::new(), tensor);
402        assert!(param.is_require_grad());
403        assert!(param.require_grad);
404
405        let param = param.valid();
406        assert!(!param.is_require_grad());
407        assert!(param.require_grad); // stateful
408
409        // Without `HasAutodiffModule`, we would need to specify the param type as well, which would be annoying:
410        // let param: Param<Tensor<TestAutodiffBackend, _>> = param.train();
411        let param = param.train();
412        assert!(param.is_require_grad());
413        assert!(param.require_grad); // stateful
414
415        let param = param.no_grad();
416        assert!(!param.is_require_grad());
417        assert!(!param.require_grad); // stateful
418
419        let param = param.valid();
420        assert!(!param.is_require_grad()); // always
421        assert!(!param.require_grad); // stateful
422
423        let param = param.train();
424        assert!(!param.is_require_grad());
425        assert!(!param.require_grad); // stateful
426    }
427}
burn_core/module/param/tensor.rs

burn_core/module/param/
tensor.rs
