burn-dispatch 0.23.0-pre.1

Backend dispatch for the Burn framework
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
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

#[cfg(feature = "autodiff")]
use alloc::boxed::Box;

#[cfg(feature = "autodiff")]
use burn_autodiff::{
    checkpoint::{
        base::Checkpointer,
        strategy::{BalancedCheckpointing, CheckpointStrategy, NoCheckpointing},
    },
    grads::Gradients,
    ops::{Backward, Ops, unary},
};
#[cfg(feature = "autodiff")]
use burn_backend::{Backend, tensor::FloatTensor};
use burn_group::{Sharding, ShardingError};

use crate::{BackendTensor, DispatchTensor, DispatchTensorKind, backends::Group};
#[cfg(feature = "autodiff")]
use crate::{DispatchAutodiffContext, GradientCheckpointingStrategy, backends::Autodiff};

impl DispatchTensor {
    /// The same value with another sharding over its device group, or why it cannot take it.
    ///
    /// A tracked autodiff tensor records the move: its gradient flows back through it unchanged,
    /// since a gradient may have any sharding.
    pub fn try_with_sharding(self, sharding: Sharding) -> Result<Self, ShardingError> {
        let DispatchTensor { kind, autodiff } = self;
        let kind = match kind {
            DispatchTensorKind::Group(tensor) => {
                DispatchTensorKind::Group(tensor.reshard(sharding)?)
            }
            #[cfg(feature = "autodiff")]
            DispatchTensorKind::Autodiff(inner) => {
                let DispatchAutodiffContext::Enabled(checkpointing) = autodiff else {
                    unreachable!("An autodiff tensor has an enabled autodiff context")
                };
                let DispatchTensorKind::Group(BackendTensor::Autodiff(tensor)) = *inner else {
                    return Err(ShardingError::NotOnGroup);
                };
                let resharded = match checkpointing {
                    GradientCheckpointingStrategy::Balanced => {
                        WithSharding::record::<BalancedCheckpointing>(tensor, sharding)
                    }
                    GradientCheckpointingStrategy::Disabled => {
                        WithSharding::record::<NoCheckpointing>(tensor, sharding)
                    }
                }?;
                DispatchTensorKind::Autodiff(Box::new(DispatchTensorKind::Group(
                    BackendTensor::Autodiff(resharded),
                )))
            }
            #[allow(unreachable_patterns)]
            _ => return Err(ShardingError::NotOnGroup),
        };
        Ok(DispatchTensor { kind, autodiff })
    }

    /// Where the tensor's shards sit, or `None` when it is not on a device group.
    pub fn sharding(&self) -> Option<Sharding> {
        let primitive = match &self.kind {
            DispatchTensorKind::Group(
                BackendTensor::Float(tensor)
                | BackendTensor::Int(tensor)
                | BackendTensor::Bool(tensor)
                | BackendTensor::Quantized(tensor),
            ) => tensor,
            #[cfg(feature = "autodiff")]
            DispatchTensorKind::Autodiff(inner) => match &**inner {
                DispatchTensorKind::Group(BackendTensor::Autodiff(tensor)) => tensor.primitive(),
                _ => return None,
            },
            #[allow(unreachable_patterns)]
            _ => return None,
        };
        Some(primitive.client.sharding(primitive))
    }
}

/// [`DispatchTensor::try_with_sharding`] on a tracked tensor, recorded so its gradient flows back.
#[cfg(feature = "autodiff")]
#[derive(Debug)]
struct WithSharding;

#[cfg(feature = "autodiff")]
impl WithSharding {
    fn record<C: CheckpointStrategy>(
        tensor: FloatTensor<Autodiff<Group>>,
        sharding: Sharding,
    ) -> Result<FloatTensor<Autodiff<Group>>, ShardingError> {
        let node = tensor.node();
        let inner = tensor.into_primitive();
        let resharded = inner.client.clone().reshard(inner, sharding)?;
        Ok(
            <WithSharding as Backward<Group, 1>>::prepare::<C>(WithSharding, [node])
                .compute_bound()
                .stateless(resharded),
        )
    }
}

#[cfg(feature = "autodiff")]
impl<B: Backend> Backward<B, 1> for WithSharding {
    type State = ();

    fn backward(self, ops: Ops<(), 1>, grads: &mut Gradients, _: &mut Checkpointer) {
        unary::<B, _>(ops.parents, ops.node, grads, |grad| grad);
    }
}

impl BackendTensor<Group> {
    fn reshard(self, sharding: Sharding) -> Result<Self, ShardingError> {
        Ok(match self {
            BackendTensor::Float(tensor) => {
                BackendTensor::Float(tensor.client.clone().reshard(tensor, sharding)?)
            }
            BackendTensor::Int(tensor) => {
                BackendTensor::Int(tensor.client.clone().reshard(tensor, sharding)?)
            }
            BackendTensor::Bool(tensor) => {
                BackendTensor::Bool(tensor.client.clone().reshard(tensor, sharding)?)
            }
            BackendTensor::Quantized(_) => return Err(ShardingError::Quantized),
            #[cfg(feature = "autodiff")]
            BackendTensor::Autodiff(_) => {
                unreachable!("An autodiff tensor is wrapped in its own kind")
            }
        })
    }
}