mirror of
https://github.com/guillaume-be/rust-bert.git
synced 2024-09-11 12:55:34 +03:00
Updated T5 for FP16 compatibility
This commit is contained in:
parent
038ac90757
commit
889f509e6c
@ -18,8 +18,6 @@ use rust_bert::resources::{RemoteResource, Resource};
|
|||||||
use rust_bert::t5::{T5ConfigResources, T5ModelResources, T5VocabResources};
|
use rust_bert::t5::{T5ConfigResources, T5ModelResources, T5VocabResources};
|
||||||
|
|
||||||
fn main() -> anyhow::Result<()> {
|
fn main() -> anyhow::Result<()> {
|
||||||
// let summarization_model = SummarizationModel::new(Default::default())?;
|
|
||||||
|
|
||||||
let config_resource =
|
let config_resource =
|
||||||
Resource::Remote(RemoteResource::from_pretrained(T5ConfigResources::T5_SMALL));
|
Resource::Remote(RemoteResource::from_pretrained(T5ConfigResources::T5_SMALL));
|
||||||
let vocab_resource =
|
let vocab_resource =
|
||||||
|
@ -224,7 +224,7 @@ impl T5Attention {
|
|||||||
scores += position_bias;
|
scores += position_bias;
|
||||||
|
|
||||||
let attention_weights = scores
|
let attention_weights = scores
|
||||||
.softmax(-1, Kind::Float)
|
.softmax(-1, scores.kind())
|
||||||
.apply_t(&self.dropout, train);
|
.apply_t(&self.dropout, train);
|
||||||
let context = self
|
let context = self
|
||||||
.unshape(attention_weights.matmul(&v), bs)
|
.unshape(attention_weights.matmul(&v), bs)
|
||||||
|
@ -18,7 +18,7 @@ use crate::t5::T5Config;
|
|||||||
use crate::RustBertError;
|
use crate::RustBertError;
|
||||||
use std::borrow::{Borrow, BorrowMut};
|
use std::borrow::{Borrow, BorrowMut};
|
||||||
use tch::nn::LinearConfig;
|
use tch::nn::LinearConfig;
|
||||||
use tch::{nn, Kind, Tensor};
|
use tch::{nn, Kind, Scalar, Tensor};
|
||||||
|
|
||||||
pub struct T5DenseReluDense {
|
pub struct T5DenseReluDense {
|
||||||
wi: nn::Linear,
|
wi: nn::Linear,
|
||||||
@ -140,6 +140,21 @@ impl T5Block {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn clamp_hidden_states(hidden_states: Tensor) -> Tensor {
|
||||||
|
if (hidden_states.kind() != Kind::Float) & bool::from(hidden_states.isinf().any()) {
|
||||||
|
let clamp_value = match hidden_states.kind() {
|
||||||
|
Kind::Half => half::f16::MAX.to_f64() - 1000.,
|
||||||
|
Kind::BFloat16 => half::bf16::MAX.to_f64() - 1000.,
|
||||||
|
_ => {
|
||||||
|
panic!("Type not supported: supported types are Float (single precision), Half and BFloat16 (half precision)");
|
||||||
|
}
|
||||||
|
};
|
||||||
|
hidden_states.clamp(Scalar::from(-clamp_value), Scalar::from(clamp_value))
|
||||||
|
} else {
|
||||||
|
hidden_states
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub fn forward_t(
|
pub fn forward_t(
|
||||||
&self,
|
&self,
|
||||||
hidden_states: &Tensor,
|
hidden_states: &Tensor,
|
||||||
@ -152,7 +167,7 @@ impl T5Block {
|
|||||||
train: bool,
|
train: bool,
|
||||||
) -> T5BlockOutput {
|
) -> T5BlockOutput {
|
||||||
let (
|
let (
|
||||||
hidden_states,
|
mut hidden_states,
|
||||||
self_attention_weights,
|
self_attention_weights,
|
||||||
self_attention_position_bias,
|
self_attention_position_bias,
|
||||||
self_attention_layer_past,
|
self_attention_layer_past,
|
||||||
@ -164,8 +179,10 @@ impl T5Block {
|
|||||||
train,
|
train,
|
||||||
);
|
);
|
||||||
|
|
||||||
|
hidden_states = T5Block::clamp_hidden_states(hidden_states);
|
||||||
|
|
||||||
let (
|
let (
|
||||||
hidden_states,
|
mut hidden_states,
|
||||||
cross_attention_weights,
|
cross_attention_weights,
|
||||||
cross_attention_position_bias,
|
cross_attention_position_bias,
|
||||||
cross_attention_layer_past,
|
cross_attention_layer_past,
|
||||||
@ -186,8 +203,12 @@ impl T5Block {
|
|||||||
(hidden_states, None, None, None)
|
(hidden_states, None, None, None)
|
||||||
};
|
};
|
||||||
|
|
||||||
|
hidden_states = T5Block::clamp_hidden_states(hidden_states);
|
||||||
|
|
||||||
layer_states = (self_attention_layer_past, cross_attention_layer_past);
|
layer_states = (self_attention_layer_past, cross_attention_layer_past);
|
||||||
let hidden_states = self.ff_layer.forward_t(&hidden_states, train);
|
let mut hidden_states = self.ff_layer.forward_t(&hidden_states, train);
|
||||||
|
|
||||||
|
hidden_states = T5Block::clamp_hidden_states(hidden_states);
|
||||||
|
|
||||||
T5BlockOutput {
|
T5BlockOutput {
|
||||||
hidden_states,
|
hidden_states,
|
||||||
@ -305,8 +326,10 @@ impl T5Stack {
|
|||||||
3 => attention_mask.unsqueeze(1),
|
3 => attention_mask.unsqueeze(1),
|
||||||
2 => {
|
2 => {
|
||||||
if self.is_decoder {
|
if self.is_decoder {
|
||||||
let seq_ids =
|
let seq_ids = Tensor::arange(
|
||||||
Tensor::arange(input_shape[1], (Kind::Float, input_embeddings.device()));
|
input_shape[1],
|
||||||
|
(input_embeddings.kind(), input_embeddings.device()),
|
||||||
|
);
|
||||||
let causal_mask = seq_ids.unsqueeze(0).unsqueeze(0).repeat(&[
|
let causal_mask = seq_ids.unsqueeze(0).unsqueeze(0).repeat(&[
|
||||||
input_shape[0],
|
input_shape[0],
|
||||||
input_shape[1],
|
input_shape[1],
|
||||||
@ -325,8 +348,10 @@ impl T5Stack {
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
let extended_attention_mask: Option<Tensor> =
|
let extended_attention_mask: Option<Tensor> = Some(
|
||||||
Some((extended_attention_mask.ones_like() - extended_attention_mask) * -1e9);
|
((extended_attention_mask.ones_like() - extended_attention_mask) * -1e4)
|
||||||
|
.to_kind(input_embeddings.kind()),
|
||||||
|
);
|
||||||
|
|
||||||
let extended_encoder_attention_mask = if self.is_decoder & encoder_hidden_states.is_some() {
|
let extended_encoder_attention_mask = if self.is_decoder & encoder_hidden_states.is_some() {
|
||||||
let encoder_hidden_states = encoder_hidden_states.as_ref().unwrap();
|
let encoder_hidden_states = encoder_hidden_states.as_ref().unwrap();
|
||||||
@ -350,7 +375,9 @@ impl T5Stack {
|
|||||||
));
|
));
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
Some((encoder_mask.ones_like() - encoder_mask) * -1e9)
|
Some(
|
||||||
|
((encoder_mask.ones_like() - encoder_mask) * -1e4).to_kind(input_embeddings.kind()),
|
||||||
|
)
|
||||||
} else {
|
} else {
|
||||||
None
|
None
|
||||||
};
|
};
|
||||||
|
@ -32,8 +32,16 @@ impl T5LayerNorm {
|
|||||||
|
|
||||||
impl Module for T5LayerNorm {
|
impl Module for T5LayerNorm {
|
||||||
fn forward(&self, x: &Tensor) -> Tensor {
|
fn forward(&self, x: &Tensor) -> Tensor {
|
||||||
let variance = x.pow(2f64).mean_dim(&[-1], true, Kind::Float);
|
let input_type = x.kind();
|
||||||
|
let variance = x
|
||||||
|
.to_kind(Kind::Float)
|
||||||
|
.pow(2.0_f64)
|
||||||
|
.mean_dim(&[-1], true, Kind::Float);
|
||||||
let x = x * (variance + self.epsilon).rsqrt();
|
let x = x * (variance + self.epsilon).rsqrt();
|
||||||
&self.weight * x
|
if input_type != Kind::Float {
|
||||||
|
(&self.weight * x).to_kind(input_type)
|
||||||
|
} else {
|
||||||
|
&self.weight * x
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
Loading…
Reference in New Issue
Block a user