huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-04-21 23:12:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Aligin the source code with main branch 2.0.4

Browse Source 
Signed-off-by: yuanwu <yuan.wu@intel.com>
This commit is contained in:
yuanwu 2024-09-24 03:06:55 +00:00
parent 0deebe7012
commit 92a1e0fbae
48 changed files with 3972 additions and 6751 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
benchmark/src/generation.rs
View File

@ -163,11 +163,8 @@ async fn prefill(
// Run prefill // Run prefill
let start_time = Instant::now(); let start_time = Instant::now();
let (_, decode_batch, _) = client.prefill(batch.clone()).await?; let (_, decode_batch, _) = client.prefill(batch.clone()).await?;
let (_, decode_batch, _) = client.decode(vec![decode_batch.clone().unwrap()]).await?;
// Get latency // Get latency
let latency = start_time.elapsed(); let latency = start_time.elapsed();
@ -183,12 +180,11 @@ async fn prefill(
}; };
Ok((step, decode_batch)) Ok((step, decode_batch))
} }
/// Run a full decode /// Run a full decode
async fn decode(batch: CachedBatch, client: &mut ShardedClient) -> Result<Decode, ClientError> { async fn decode(batch: CachedBatch, client: &mut ShardedClient) -> Result<Decode, ClientError> {
let mut decode_length = 1; // 1 decode step was already scheduled in prefill with speculative scheduling let mut decode_length = 0;
let batch_size = batch.size; let batch_size = batch.size;
let start_time = Instant::now(); let start_time = Instant::now();

2
benchmark/src/main.rs
View File

@ -51,7 +51,7 @@ struct Args {
runs: usize, runs: usize,
/// Number of warmup cycles /// Number of warmup cycles
#[clap(default_value = "3", short, long, env)] #[clap(default_value = "1", short, long, env)]
warmups: usize, warmups: usize,
/// The location of the grpc socket. This benchmark tool bypasses the router /// The location of the grpc socket. This benchmark tool bypasses the router

1
router/client/Cargo.toml
View File

@ -9,7 +9,6 @@ homepage.workspace = true
futures = "^0.3" futures = "^0.3"
grpc-metadata = { path = "../grpc-metadata" } grpc-metadata = { path = "../grpc-metadata" }
prost = "^0.12" prost = "^0.12"
rand = "0.8.5"
thiserror = "^1.0" thiserror = "^1.0"
tokio = { version = "^1.32", features = ["sync"] } tokio = { version = "^1.32", features = ["sync"] }
tonic = "^0.10" tonic = "^0.10"

296
router/client/src/client.rs
View File

@ -1,13 +1,9 @@
/// Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.
/// Single shard Client /// Single shard Client
use crate::pb::generate::v2::text_generation_service_client::TextGenerationServiceClient; use crate::pb::generate::v2::text_generation_service_client::TextGenerationServiceClient;
use crate::pb::generate::v2::*; use crate::pb::generate::v2::*;
use crate::Result; use crate::Result;
use std::env;
use rand::{distributions::Uniform, Rng};
use grpc_metadata::InjectTelemetryContext; use grpc_metadata::InjectTelemetryContext;
use std::cmp; use std::cmp::min;
use std::time::Duration; use std::time::Duration;
use tonic::transport::{Channel, Uri}; use tonic::transport::{Channel, Uri};
use tracing::instrument; use tracing::instrument;
@ -110,59 +106,20 @@ impl Client {
max_prefill_tokens: u32, max_prefill_tokens: u32,
max_total_tokens: u32, max_total_tokens: u32,
max_batch_size: Option<usize>, max_batch_size: Option<usize>,
model_id: &str
) -> Result<Option<u32>> { ) -> Result<Option<u32>> {
let warmup_enabled: bool = env::var("WARMUP_ENABLED").ok().map_or(true, |value| value.to_lowercase() == "true");
if !warmup_enabled {
return Ok(None);
}
let read_env_var = |key: &str, default: u32| -> u32 {
env::var(key).ok().map_or(default, |value| value.parse::<u32>().unwrap())
};
// get all possible batch sizes
let decode_bucket_size: u32 = read_env_var("BATCH_BUCKET_SIZE", 8);
let max_decode_batch_size: u32 = match max_batch_size {
Some(max_batch_size) => max_batch_size as u32,
None => decode_bucket_size
};
let decode_batch_sizes: Vec<u32> = (decode_bucket_size..max_decode_batch_size+1).step_by(decode_bucket_size as usize).collect();
let prefill_bucket_size: u32 = read_env_var("PREFILL_BATCH_BUCKET_SIZE", 4);
let mut max_prefill_batch_size: u32 = max_prefill_tokens / max_input_length;
max_prefill_batch_size = cmp::min(max_prefill_batch_size, max_decode_batch_size);
let prefill_batch_sizes: Vec<u32> = (prefill_bucket_size..max_prefill_batch_size+1).step_by(prefill_bucket_size as usize).collect();
// get all possible sequence lengths for prefill
let seq_bucket_size: u32 = read_env_var("PAD_SEQUENCE_TO_MULTIPLE_OF", 128);
let mut seq_lengths: Vec<u32> = (seq_bucket_size..max_input_length+1).step_by(seq_bucket_size as usize).collect();
if let Some(&last) = seq_lengths.last() {
if last < max_input_length {
seq_lengths.push(max_input_length);
}
}
// execute batch for each combination of batch size and sequence length
let mut shapes: Vec<(u32, u32)> = Vec::with_capacity(prefill_batch_sizes.len() * seq_lengths.len());
for batch_size in &prefill_batch_sizes {
for seq_length in &seq_lengths {
shapes.push((*batch_size, *seq_length));
}
}
let mut batch_counter: u64 = 0;
let mut request_counter: u64 = 0;
if model_id.contains("llava") {
let mut n_tokens = 0; let mut n_tokens = 0;
let mut requests = Vec::new(); let mut requests = Vec::new();
// Create requests // Create requests
while n_tokens < max_prefill_tokens { while n_tokens < max_prefill_tokens {
let truncate = cmp::min(max_input_length, max_prefill_tokens - n_tokens); let truncate = min(max_input_length, max_prefill_tokens - n_tokens);
let mut inputs = String::new(); let mut inputs = String::new();
inputs.push_str("![](data:image/jpeg;base64,iVBORw0KGgoAAAANSUhEUgAAABQAAAAUCAIAAAAC64paAAABg2lDQ1BJQ0MgcHJvZmlsZQAAKJF9kT1Iw0AcxV/TSotUROxQxCFDdbKLijjWKhShQqgVWnUwufQLmrQkKS6OgmvBwY/FqoOLs64OroIg+AHi7OCk6CIl/i8ptIjx4Lgf7+497t4BQqvKNDOQADTdMjKppJjLr4rBVwQQwhAERGVm1uckKQ3P8XUPH1/v4jzL+9yfY0AtmAzwicQJVjcs4g3imU2rznmfOMLKskp8Tjxh0AWJH7muuPzGueSwwDMjRjYzTxwhFks9rPQwKxsa8TRxTNV0yhdyLquctzhr1Qbr3JO/MFzQV5a5TnMUKSxiCRJEKGiggiosxGnVSTGRof2kh3/E8UvkUshVASPHAmrQIDt+8D/43a1ZnJp0k8JJoO/Ftj/GgOAu0G7a9vexbbdPAP8zcKV3/bUWMPtJerOrxY6AwW3g4rqrKXvA5Q4QfarLhuxIfppCsQi8n9E35YHhW6B/ze2ts4/TByBLXaVvgINDYLxE2ese7w719vbvmU5/PycecohsjayNAAAACXBIWXMAAC4jAAAuIwF4pT92AAAAB3RJTUUH6AQIEQMnlTSSjwAAABl0RVh0Q29tbWVudABDcmVhdGVkIHdpdGggR0lNUFeBDhcAAAASSURBVDjLY2AYBaNgFIyCoQsABMQAAeRw1DoAAAAASUVORK5CYII=)");
inputs.push_str(&"_test ".to_string().repeat(max_input_length as usize)); inputs.push_str(&"_test ".to_string().repeat(max_input_length as usize));
if n_tokens == 0 {
// 1 request is enough to test vision heads.
// Sending images on other queries messes up easily with truncation.
inputs.push_str("![](data:image/jpeg;base64,iVBORw0KGgoAAAANSUhEUgAAABQAAAAUCAIAAAAC64paAAABg2lDQ1BJQ0MgcHJvZmlsZQAAKJF9kT1Iw0AcxV/TSotUROxQxCFDdbKLijjWKhShQqgVWnUwufQLmrQkKS6OgmvBwY/FqoOLs64OroIg+AHi7OCk6CIl/i8ptIjx4Lgf7+497t4BQqvKNDOQADTdMjKppJjLr4rBVwQQwhAERGVm1uckKQ3P8XUPH1/v4jzL+9yfY0AtmAzwicQJVjcs4g3imU2rznmfOMLKskp8Tjxh0AWJH7muuPzGueSwwDMjRjYzTxwhFks9rPQwKxsa8TRxTNV0yhdyLquctzhr1Qbr3JO/MFzQV5a5TnMUKSxiCRJEKGiggiosxGnVSTGRof2kh3/E8UvkUshVASPHAmrQIDt+8D/43a1ZnJp0k8JJoO/Ftj/GgOAu0G7a9vexbbdPAP8zcKV3/bUWMPtJerOrxY6AwW3g4rqrKXvA5Q4QfarLhuxIfppCsQi8n9E35YHhW6B/ze2ts4/TByBLXaVvgINDYLxE2ese7w719vbvmU5/PycecohsjayNAAAACXBIWXMAAC4jAAAuIwF4pT92AAAAB3RJTUUH6AQIEQMnlTSSjwAAABl0RVh0Q29tbWVudABDcmVhdGVkIHdpdGggR0lNUFeBDhcAAAASSURBVDjLY2AYBaNgFIyCoQsABMQAAeRw1DoAAAAASUVORK5CYII=)");
}
requests.push(Request { requests.push(Request {
id: 0, id: 0,
@ -199,16 +156,15 @@ impl Client {
} }
} }
let mut batches = Vec::new(); let batch = Batch {
batches.push(Batch {
id: 0, id: 0,
size: requests.len() as u32, size: requests.len() as u32,
requests, requests,
max_tokens: 0, max_tokens: 0,
}); };
let request = tonic::Request::new(WarmupRequest { let request = tonic::Request::new(WarmupRequest {
batches, batch: Some(batch),
max_input_length, max_input_length,
max_prefill_tokens, max_prefill_tokens,
max_total_tokens, max_total_tokens,
@ -217,238 +173,6 @@ impl Client {
let response = self.stub.warmup(request).await?.into_inner(); let response = self.stub.warmup(request).await?.into_inner();
Ok(response.max_supported_total_tokens) Ok(response.max_supported_total_tokens)
} }
else {
for shape in shapes.iter() {
let (batch_size, seq_length) = shape;
let mut batches: Vec<Batch> = vec![
self.create_warmup_batch(
*shape,
&mut batch_counter,
&mut request_counter,
max_input_length,
max_total_tokens,
seq_bucket_size,
false,
None,
)
];
// if possible, create second batch in order to trigger concatenate operation
if *batch_size < max_decode_batch_size {
batches.push(
self.create_warmup_batch(
(1, *seq_length),
&mut batch_counter,
&mut request_counter,
max_input_length,
max_total_tokens,
seq_bucket_size,
false,
None,
)
);
}
let request = tonic::Request::new(WarmupRequest {
batches,
max_input_length,
max_prefill_tokens,
max_total_tokens,
}).inject_context();
let _response = self.stub.warmup(request).await?.into_inner();
}
// send batches to warmup all possible decode shapes
if decode_batch_sizes.len() > 1 {
let steps_per_bucket: u32 = if decode_bucket_size <= max_prefill_batch_size {
decode_bucket_size
} else {
decode_bucket_size.div_ceil(max_prefill_batch_size)
};
let max_new_tokens: u32 = 2 * decode_batch_sizes.len() as u32 * steps_per_bucket;
let mut requests_send: u32 = cmp::min(max_prefill_batch_size, decode_bucket_size);
let mut batches: Vec<Batch> = vec![
self.create_warmup_batch(
(requests_send, seq_bucket_size),
&mut batch_counter,
&mut request_counter,
max_input_length,
max_total_tokens,
seq_bucket_size,
false,
Some(max_new_tokens),
)
];
let get_current_decode_batch_size = |num: u32| -> u32 {
decode_batch_sizes.iter()
.filter(|&&x| x >= num)
.min()
.copied()
.unwrap()
};
let mut current_decode_batch_size: u32 = get_current_decode_batch_size(requests_send);
while current_decode_batch_size < max_decode_batch_size {
let distance_to_next_bucket = current_decode_batch_size + decode_bucket_size - requests_send;
let num_requests: u32 = cmp::min(distance_to_next_bucket, max_prefill_batch_size);
batches.push(
self.create_warmup_batch(
(num_requests, seq_bucket_size),
&mut batch_counter,
&mut request_counter,
max_input_length,
max_total_tokens,
seq_bucket_size,
false,
Some(max_new_tokens),
)
);
requests_send += num_requests;
current_decode_batch_size = get_current_decode_batch_size(requests_send);
}
let request = tonic::Request::new(WarmupRequest {
batches,
max_input_length,
max_prefill_tokens,
max_total_tokens,
}).inject_context();
let _response = self.stub.warmup(request).await?.into_inner();
}
// send batches with default params to warm up Greedy search
let mut greedy_shapes: Vec<(u32, u32)> = Vec::with_capacity(prefill_batch_sizes.len());
for batch_size in &prefill_batch_sizes {
greedy_shapes.push((*batch_size, seq_bucket_size.clone()));
}
for greedy_shape in greedy_shapes.iter() {
let batches: Vec<Batch> = vec![
self.create_warmup_batch(
*greedy_shape,
&mut batch_counter,
&mut request_counter,
max_input_length,
max_total_tokens,
seq_bucket_size,
true,
None,
)
];
let request = tonic::Request::new(WarmupRequest {
batches,
max_input_length,
max_prefill_tokens,
max_total_tokens,
}).inject_context();
let _response = self.stub.warmup(request).await?.into_inner();
}
Ok(None) // No support for maximum total tokens
}
}
#[instrument(skip_all)]
fn create_warmup_batch(
&mut self,
shape: (u32, u32),
batch_counter: &mut u64,
request_counter: &mut u64,
max_input_length: u32,
max_total_tokens: u32,
seq_bucket_size: u32,
default_params: bool,
max_new_tokens: Option<u32>,
) -> Batch {
*batch_counter += 1;
let (batch_size, input_length) = shape;
let mut requests = Vec::new();
for _ in 0..batch_size {
*request_counter += 1;
let req_params = if default_params {
Some(NextTokenChooserParameters {
temperature: 1.0,
top_k: 0,
top_p: 1.0,
typical_p: 1.0,
do_sample: false,
seed: 0,
repetition_penalty: 1.0,
frequency_penalty: 0.0,
watermark: false,
grammar: String::new(),
grammar_type: GrammarType::None as i32,
})
} else {
Some(NextTokenChooserParameters {
temperature: 0.9,
top_k: 10,
top_p: 0.9,
typical_p: 0.9,
do_sample: true,
seed: 0,
repetition_penalty: 1.2,
frequency_penalty: 0.1,
watermark: false,
grammar: String::new(),
grammar_type: GrammarType::None as i32,
})
};
requests.push(Request {
id: *request_counter,
inputs: self.get_random_input(input_length, seq_bucket_size),
truncate: max_input_length,
parameters: req_params,
stopping_parameters: Some(StoppingCriteriaParameters {
max_new_tokens: cmp::min(max_new_tokens.unwrap_or(10), max_total_tokens - max_input_length),
stop_sequences: vec![],
ignore_eos_token: true,
}),
prefill_logprobs: false,
top_n_tokens: 0,
});
}
Batch {
id: *batch_counter,
size: requests.len() as u32,
requests,
max_tokens: max_total_tokens,
}
}
#[instrument(skip_all)]
fn get_random_input(
&mut self,
input_length: u32,
seq_bucket_size: u32,
) -> String {
let skip_tokenizer_in_tgi: bool = env::var("SKIP_TOKENIZER_IN_TGI")
.ok()
.map_or(false, |value| value.to_lowercase() == "true");
if skip_tokenizer_in_tgi {
// generate random tokens
let mut rng = rand::thread_rng();
let range = Uniform::new(2, 8192);
let tokens = if input_length % seq_bucket_size == 0 {
input_length - seq_bucket_size / 2
} else {
input_length - (input_length % seq_bucket_size) / 2
};
(0..tokens)
.map(|_| rng.sample(&range).to_string())
.collect::<Vec<String>>()
.join(", ")
} else {
// repeat test string to get expected input shape
let mut bucket_id = input_length / seq_bucket_size;
if input_length % seq_bucket_size != 0 {
bucket_id += 1
}
let repeats = cmp::max(1, (bucket_id - 1) * seq_bucket_size / 2);
"_test ".to_string().repeat(repeats as usize)
}
}
/// Generate one token for each request in the given batch /// Generate one token for each request in the given batch
/// ///

4
router/client/src/sharded_client.rs
View File

@ -1,5 +1,3 @@
/// Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.
use crate::client::{DecodeTimings, PrefillTimings}; use crate::client::{DecodeTimings, PrefillTimings};
/// Multi shard Client /// Multi shard Client
use crate::{Batch, CachedBatch, Client, Generation, HealthResponse, ShardInfo}; use crate::{Batch, CachedBatch, Client, Generation, HealthResponse, ShardInfo};
@ -100,7 +98,6 @@ impl ShardedClient {
max_prefill_tokens: u32, max_prefill_tokens: u32,
max_total_tokens: u32, max_total_tokens: u32,
max_batch_size: Option<usize>, max_batch_size: Option<usize>,
model_id: &str,
) -> Result<Option<u32>> { ) -> Result<Option<u32>> {
let futures: Vec<_> = self let futures: Vec<_> = self
.clients .clients
@ -111,7 +108,6 @@ impl ShardedClient {
max_prefill_tokens, max_prefill_tokens,
max_total_tokens, max_total_tokens,
max_batch_size, max_batch_size,
model_id
)) ))
}) })
.collect(); .collect();

13
router/src/infer.rs
View File

@ -1,5 +1,3 @@
/// Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.
/// Batching and inference logic /// Batching and inference logic
use crate::validation::{Validation, ValidationError}; use crate::validation::{Validation, ValidationError};
use crate::{ use crate::{
@ -65,22 +63,13 @@ impl Infer {
max_batch_size: Option<usize>, max_batch_size: Option<usize>,
max_concurrent_requests: usize, max_concurrent_requests: usize,
requires_padding: bool, requires_padding: bool,
max_input_length: u32,
max_total_tokens: u32,
window_size: Option<u32>, window_size: Option<u32>,
speculate: u32, speculate: u32,
generation_health: Arc<AtomicBool>, generation_health: Arc<AtomicBool>,
tokenizer_config: HubTokenizerConfig, tokenizer_config: HubTokenizerConfig,
) -> Self { ) -> Self {
// Infer shared state // Infer shared state
let queue = Queue::new( let queue = Queue::new(requires_padding, 16, window_size, speculate);
requires_padding,
max_input_length,
max_total_tokens,
16,
window_size,
speculate
);
let shared = Arc::new(Shared { let shared = Arc::new(Shared {
batching_task: Notify::new(), batching_task: Notify::new(),
}); });

33
router/src/main.rs
View File

@ -1,5 +1,3 @@
/// Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.
use axum::http::HeaderValue; use axum::http::HeaderValue;
use clap::Parser; use clap::Parser;
use hf_hub::api::tokio::{Api, ApiBuilder, ApiRepo}; use hf_hub::api::tokio::{Api, ApiBuilder, ApiRepo};
@ -10,7 +8,6 @@ use opentelemetry::sdk::trace::Sampler;
use opentelemetry::sdk::Resource; use opentelemetry::sdk::Resource;
use opentelemetry::{global, KeyValue}; use opentelemetry::{global, KeyValue};
use opentelemetry_otlp::WithExportConfig; use opentelemetry_otlp::WithExportConfig;
use std::env;
use std::fs::File; use std::fs::File;
use std::io::BufReader; use std::io::BufReader;
use std::net::{IpAddr, Ipv4Addr, SocketAddr}; use std::net::{IpAddr, Ipv4Addr, SocketAddr};
@ -148,25 +145,6 @@ async fn main() -> Result<(), RouterError> {
} }
} }
let (max_batch_size, max_batch_total_tokens) = match (max_batch_size, max_batch_total_tokens) {
(Some(_max_batch_size), Some(_max_batch_total_tokens)) => {
if (_max_batch_total_tokens as usize / max_total_tokens) != _max_batch_size {
tracing::warn!("max_batch_size was set to {_max_batch_size} while max_batch_total_tokens to {_max_batch_total_tokens}");
tracing::warn!("These values are not match, so max_batch_size will be preferred");
(Some(_max_batch_size), Some((_max_batch_size * max_total_tokens) as u32))
} else {
(Some(_max_batch_size), Some(_max_batch_total_tokens))
}
},
(Some(_max_batch_size), None) => (
Some(_max_batch_size), Some((_max_batch_size * max_total_tokens) as u32)
),
(None, Some(_max_batch_total_tokens)) => (
Some(_max_batch_total_tokens as usize / max_total_tokens), Some(_max_batch_total_tokens)
),
(None, None) => (None, None),
};
// CORS allowed origins // CORS allowed origins
// map to go inside the option and then map to parse from String to HeaderValue // map to go inside the option and then map to parse from String to HeaderValue
// Finally, convert to AllowOrigin // Finally, convert to AllowOrigin
@ -228,9 +206,6 @@ async fn main() -> Result<(), RouterError> {
}; };
// Load tokenizer and model info // Load tokenizer and model info
let skip_tokenizer_in_tgi = env::var("SKIP_TOKENIZER_IN_TGI")
.ok()
.map_or(false, |value| value.to_lowercase() == "true");
let (tokenizer_filename, config_filename, tokenizer_config_filename, model_info) = match api { let (tokenizer_filename, config_filename, tokenizer_config_filename, model_info) = match api {
Type::None => ( Type::None => (
Some(local_path.join("tokenizer.json")), Some(local_path.join("tokenizer.json")),
@ -279,11 +254,8 @@ async fn main() -> Result<(), RouterError> {
) )
} }
}; };
let tokenizer: Option<Tokenizer> = if skip_tokenizer_in_tgi { let tokenizer: Option<Tokenizer> =
None tokenizer_filename.and_then(|filename| Tokenizer::from_file(filename).ok());
} else {
tokenizer_filename.and_then(|filename| Tokenizer::from_file(filename).ok())
};
let config: Option<Config> = config_filename.and_then(|filename| { let config: Option<Config> = config_filename.and_then(|filename| {
std::fs::read_to_string(filename) std::fs::read_to_string(filename)
.ok() .ok()
@ -349,7 +321,6 @@ async fn main() -> Result<(), RouterError> {
max_batch_prefill_tokens, max_batch_prefill_tokens,
max_total_tokens as u32, max_total_tokens as u32,
max_batch_size, max_batch_size,
&model_info.model_id
) )
.await .await
.map_err(RouterError::Warmup)? .map_err(RouterError::Warmup)?

211
router/src/queue.rs
View File

@ -1,14 +1,9 @@
/// Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.
use crate::infer::InferError; use crate::infer::InferError;
use crate::infer::InferStreamResponse; use crate::infer::InferStreamResponse;
use crate::validation::ValidGenerateRequest; use crate::validation::ValidGenerateRequest;
use nohash_hasher::{BuildNoHashHasher, IntMap}; use nohash_hasher::{BuildNoHashHasher, IntMap};
use std::cmp::min; use std::cmp::min;
use std::cmp::{Eq, Ord, PartialEq, PartialOrd}; use std::collections::VecDeque;
use std::collections::BinaryHeap;
use std::env;
use std::time::Duration;
use text_generation_client::{Batch, Request}; use text_generation_client::{Batch, Request};
use tokio::sync::{mpsc, oneshot}; use tokio::sync::{mpsc, oneshot};
use tokio::time::Instant; use tokio::time::Instant;
@ -41,8 +36,6 @@ pub(crate) struct Queue {
impl Queue { impl Queue {
pub(crate) fn new( pub(crate) fn new(
requires_padding: bool, requires_padding: bool,
max_input_length: u32,
max_total_tokens: u32,
block_size: u32, block_size: u32,
window_size: Option<u32>, window_size: Option<u32>,
speculate: u32, speculate: u32,
@ -53,8 +46,6 @@ impl Queue {
// Launch background queue task // Launch background queue task
tokio::spawn(queue_task( tokio::spawn(queue_task(
requires_padding, requires_padding,
max_input_length,
max_total_tokens,
block_size, block_size,
window_size, window_size,
speculate, speculate,
@ -106,21 +97,12 @@ impl Queue {
// Background task responsible of the queue state // Background task responsible of the queue state
async fn queue_task( async fn queue_task(
requires_padding: bool, requires_padding: bool,
max_input_length: u32,
max_total_tokens: u32,
block_size: u32, block_size: u32,
window_size: Option<u32>, window_size: Option<u32>,
speculate: u32, speculate: u32,
mut receiver: mpsc::UnboundedReceiver<QueueCommand>, mut receiver: mpsc::UnboundedReceiver<QueueCommand>,
) { ) {
let mut state = State::new( let mut state = State::new(requires_padding, block_size, window_size, speculate);
requires_padding,
max_input_length,
max_total_tokens,
block_size,
window_size,
speculate
);
while let Some(cmd) = receiver.recv().await { while let Some(cmd) = receiver.recv().await {
match cmd { match cmd {
@ -145,104 +127,11 @@ async fn queue_task(
} }
} }
#[derive(Debug)]
struct IdentifiableEntry(u64, Entry);
impl Eq for IdentifiableEntry {}
impl PartialEq for IdentifiableEntry {
fn eq(&self, other: &Self) -> bool {
self.0 == other.0
}
}
impl Ord for IdentifiableEntry {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
let ordering = match self
.1
.request
.input_length
.cmp(&other.1.request.input_length)
{
std::cmp::Ordering::Equal => self.0.cmp(&other.0),
any => any,
};
// inverse to get min heap
return ordering.reverse();
}
}
impl PartialOrd for IdentifiableEntry {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
#[derive(Debug)]
struct QueueImpl {
regular_entries: BinaryHeap<IdentifiableEntry>,
overdue_entries: BinaryHeap<IdentifiableEntry>,
overdue_threshold: Duration,
}
impl QueueImpl {
fn new(capacity: usize, overdue_threshold: Duration) -> Self {
Self {
regular_entries: BinaryHeap::with_capacity(capacity),
overdue_entries: BinaryHeap::with_capacity(capacity),
overdue_threshold,
}
}
fn update(&mut self) {
if self.regular_entries.is_empty() {
return;
}
let mut left = BinaryHeap::with_capacity(self.regular_entries.capacity());
for entry in self.regular_entries.drain() {
if entry.1.queue_time.elapsed() > self.overdue_threshold {
self.overdue_entries.push(entry);
} else {
left.push(entry);
}
}
self.regular_entries = left;
}
fn push(&mut self, entry: IdentifiableEntry) {
if entry.1.queue_time.elapsed() > self.overdue_threshold {
self.overdue_entries.push(entry);
} else {
self.regular_entries.push(entry);
}
}
fn pop(&mut self) -> Option<IdentifiableEntry> {
if !self.overdue_entries.is_empty() {
self.overdue_entries.pop()
} else {
self.regular_entries.pop()
}
}
fn is_empty(&self) -> bool {
self.regular_entries.is_empty() && self.overdue_entries.is_empty()
}
fn len(&self) -> usize {
self.regular_entries.len() + self.overdue_entries.len()
}
}
/// Queue State /// Queue State
#[derive(Debug)] #[derive(Debug)]
struct State { struct State {
/// Queue entries /// Queue entries organized in a Vec
entries: QueueImpl, entries: VecDeque<(u64, Entry)>,
/// Id of the next entry /// Id of the next entry
next_id: u64, next_id: u64,
@ -253,12 +142,6 @@ struct State {
/// Whether the model is using padding /// Whether the model is using padding
requires_padding: bool, requires_padding: bool,
/// Maximum input length, required for padding scenario
max_input_length: u32,
/// Maximum input and output length, required for padding scenario
max_total_tokens: u32,
/// Paged Attention block size /// Paged Attention block size
block_size: u32, block_size: u32,
@ -272,25 +155,15 @@ struct State {
impl State { impl State {
fn new( fn new(
requires_padding: bool, requires_padding: bool,
max_input_length: u32,
max_total_tokens: u32,
block_size: u32, block_size: u32,
window_size: Option<u32>, window_size: Option<u32>,
speculate: u32, speculate: u32,
) -> Self { ) -> Self {
let default_threshold: u64 = 120;
let threshold: u64 = match env::var("QUEUE_THRESHOLD_MS") {
Ok(val) => val.parse().unwrap_or(default_threshold),
Err(_) => default_threshold,
};
Self { Self {
entries: QueueImpl::new(128, Duration::from_millis(threshold)), entries: VecDeque::with_capacity(128),
next_id: 0, next_id: 0,
next_batch_id: 0, next_batch_id: 0,
requires_padding, requires_padding,
max_input_length,
max_total_tokens,
block_size, block_size,
window_size, window_size,
speculate, speculate,
@ -304,7 +177,7 @@ impl State {
entry.temp_span = Some(queue_span); entry.temp_span = Some(queue_span);
// Push entry in the queue // Push entry in the queue
self.entries.push(IdentifiableEntry(self.next_id, entry)); self.entries.push_back((self.next_id, entry));
self.next_id += 1; self.next_id += 1;
} }
@ -329,7 +202,9 @@ impl State {
} }
} }
self.entries.update(); // Pad prefill_token_budget to be a multiple of block size
let prefill_token_budget =
((prefill_token_budget + self.block_size - 1) / self.block_size) * self.block_size;
// Create span for this batch to add context to inference calls // Create span for this batch to add context to inference calls
let next_batch_span = info_span!(parent: None, "batch", batch_size = tracing::field::Empty); let next_batch_span = info_span!(parent: None, "batch", batch_size = tracing::field::Empty);
@ -339,11 +214,12 @@ impl State {
let mut batch_entries = let mut batch_entries =
IntMap::with_capacity_and_hasher(self.entries.len(), BuildNoHashHasher::default()); IntMap::with_capacity_and_hasher(self.entries.len(), BuildNoHashHasher::default());
let mut max_input_length = 0;
let mut prefill_tokens: u32 = 0; let mut prefill_tokens: u32 = 0;
let mut decode_tokens: u32 = 0; let mut decode_tokens: u32 = 0;
// Pop entries starting from the front of the queue // Pop entries starting from the front of the queue
while let Some(IdentifiableEntry(id, mut entry)) = self.entries.pop() { while let Some((id, mut entry)) = self.entries.pop_front() {
// Filter entries where the response receiver was dropped (== entries where the request // Filter entries where the response receiver was dropped (== entries where the request
// was dropped by the client) // was dropped by the client)
if entry.response_tx.is_closed() { if entry.response_tx.is_closed() {
@ -355,7 +231,8 @@ impl State {
if self.requires_padding { if self.requires_padding {
// We pad to max input length in the Python shards // We pad to max input length in the Python shards
// We need to take these padding tokens into the equation // We need to take these padding tokens into the equation
prefill_tokens = (batch_requests.len() + 1) as u32 * self.max_input_length; max_input_length = max_input_length.max(entry.request.input_length);
prefill_tokens = (batch_requests.len() + 1) as u32 * max_input_length
} else { } else {
// pad to block size // pad to block size
prefill_tokens += ((entry.request.input_length + self.block_size - 1) prefill_tokens += ((entry.request.input_length + self.block_size - 1)
@ -364,9 +241,7 @@ impl State {
} }
if self.requires_padding { if self.requires_padding {
// We pad to max total tokens in the Python shards decode_tokens += entry.request.stopping_parameters.max_new_tokens;
// We need to take these padding tokens into the equation
decode_tokens = (batch_requests.len() + 1) as u32 * (self.max_total_tokens - self.max_input_length);
} else { } else {
let max_new_tokens = match self.window_size { let max_new_tokens = match self.window_size {
None => entry.request.stopping_parameters.max_new_tokens, None => entry.request.stopping_parameters.max_new_tokens,
@ -387,7 +262,7 @@ impl State {
// Entry is over budget // Entry is over budget
// Add it back to the front // Add it back to the front
tracing::debug!("Over budget: prefill_tokens={prefill_tokens} > {prefill_token_budget} || {prefill_tokens} + {decode_tokens} + {} > {token_budget}", self.speculate); tracing::debug!("Over budget: prefill_tokens={prefill_tokens} > {prefill_token_budget} || {prefill_tokens} + {decode_tokens} + {} > {token_budget}", self.speculate);
self.entries.push(IdentifiableEntry(id, entry)); self.entries.push_front((id, entry));
break; break;
} }
@ -434,7 +309,7 @@ impl State {
for r in batch_requests.into_iter().rev() { for r in batch_requests.into_iter().rev() {
let id = r.id; let id = r.id;
let entry = batch_entries.remove(&id).unwrap(); let entry = batch_entries.remove(&id).unwrap();
self.entries.push(IdentifiableEntry(id, entry)); self.entries.push_front((id, entry));
} }
return None; return None;
@ -483,18 +358,6 @@ mod tests {
}; };
use tracing::info_span; use tracing::info_span;
fn default_queue() -> Queue {
Queue::new(
true, 1, 2, 1, None, 0
)
}
fn default_state() -> State {
State::new(
true, 1, 2, 1, None, 0
)
}
fn default_entry() -> ( fn default_entry() -> (
Entry, Entry,
mpsc::UnboundedReceiver<Result<InferStreamResponse, InferError>>, mpsc::UnboundedReceiver<Result<InferStreamResponse, InferError>>,
@ -538,7 +401,7 @@ mod tests {
#[test] #[test]
fn test_append() { fn test_append() {
let mut state = default_state(); let mut state = State::new(false, 1, None, 0);
let (entry, _guard) = default_entry(); let (entry, _guard) = default_entry();
assert_eq!(state.next_id, 0); assert_eq!(state.next_id, 0);
@ -548,13 +411,13 @@ mod tests {
assert_eq!(state.next_id, 1); assert_eq!(state.next_id, 1);
assert_eq!(state.entries.len(), 1); assert_eq!(state.entries.len(), 1);
let id = state.entries.pop().unwrap().0; let (id, _) = state.entries.remove(0).unwrap();
assert_eq!(id, 0); assert_eq!(id, 0);
} }
#[test] #[test]
fn test_next_batch_empty() { fn test_next_batch_empty() {
let mut state = default_state(); let mut state = State::new(false, 1, None, 0);
assert!(state.next_batch(None, None, 1, 1).is_none()); assert!(state.next_batch(None, None, 1, 1).is_none());
assert!(state.next_batch(Some(1), None, 1, 1).is_none()); assert!(state.next_batch(Some(1), None, 1, 1).is_none());
@ -562,13 +425,13 @@ mod tests {
#[test] #[test]
fn test_next_batch_min_size() { fn test_next_batch_min_size() {
let mut state = default_state(); let mut state = State::new(false, 1, None, 0);
let (entry1, _guard1) = default_entry(); let (entry1, _guard1) = default_entry();
let (entry2, _guard2) = default_entry(); let (entry2, _guard2) = default_entry();
state.append(entry1); state.append(entry1);
state.append(entry2); state.append(entry2);
let (entries, batch, _) = state.next_batch(None, None, 2, 4).unwrap(); let (entries, batch, _) = state.next_batch(None, None, 2, 2).unwrap();
assert_eq!(entries.len(), 2); assert_eq!(entries.len(), 2);
assert!(entries.contains_key(&0)); assert!(entries.contains_key(&0));
assert!(entries.contains_key(&1)); assert!(entries.contains_key(&1));
@ -588,13 +451,13 @@ mod tests {
assert_eq!(state.next_id, 3); assert_eq!(state.next_id, 3);
assert_eq!(state.entries.len(), 1); assert_eq!(state.entries.len(), 1);
let IdentifiableEntry(id, _) = state.entries.pop().unwrap(); let (id, _) = state.entries.remove(0).unwrap();
assert_eq!(id, 2); assert_eq!(id, 2);
} }
#[test] #[test]
fn test_next_batch_max_size() { fn test_next_batch_max_size() {
let mut state = default_state(); let mut state = State::new(false, 1, None, 0);
let (entry1, _guard1) = default_entry(); let (entry1, _guard1) = default_entry();
let (entry2, _guard2) = default_entry(); let (entry2, _guard2) = default_entry();
state.append(entry1); state.append(entry1);
@ -614,13 +477,13 @@ mod tests {
#[test] #[test]
fn test_next_batch_token_budget() { fn test_next_batch_token_budget() {
let mut state = default_state(); let mut state = State::new(false, 1, None, 0);
let (entry1, _guard1) = default_entry(); let (entry1, _guard1) = default_entry();
let (entry2, _guard2) = default_entry(); let (entry2, _guard2) = default_entry();
state.append(entry1); state.append(entry1);
state.append(entry2); state.append(entry2);
let (entries, batch, _) = state.next_batch(None, None, 1, 2).unwrap(); let (entries, batch, _) = state.next_batch(None, None, 1, 1).unwrap();
assert_eq!(entries.len(), 1); assert_eq!(entries.len(), 1);
assert!(entries.contains_key(&0)); assert!(entries.contains_key(&0));
assert_eq!(batch.id, 0); assert_eq!(batch.id, 0);
@ -633,7 +496,7 @@ mod tests {
let (entry3, _guard3) = default_entry(); let (entry3, _guard3) = default_entry();
state.append(entry3); state.append(entry3);
let (entries, batch, _) = state.next_batch(None, None, 3, 6).unwrap(); let (entries, batch, _) = state.next_batch(None, None, 3, 3).unwrap();
assert_eq!(entries.len(), 2); assert_eq!(entries.len(), 2);
assert!(entries.contains_key(&1)); assert!(entries.contains_key(&1));
assert!(entries.contains_key(&2)); assert!(entries.contains_key(&2));
@ -647,14 +510,14 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_queue_append() { async fn test_queue_append() {
let queue = default_queue(); let queue = Queue::new(false, 1, None, 0);
let (entry, _guard) = default_entry(); let (entry, _guard) = default_entry();
queue.append(entry); queue.append(entry);
} }
#[tokio::test] #[tokio::test]
async fn test_queue_next_batch_empty() { async fn test_queue_next_batch_empty() {
let queue = default_queue(); let queue = Queue::new(false, 1, None, 0);
assert!(queue.next_batch(None, None, 1, 1).await.is_none()); assert!(queue.next_batch(None, None, 1, 1).await.is_none());
assert!(queue.next_batch(Some(1), None, 1, 1).await.is_none()); assert!(queue.next_batch(Some(1), None, 1, 1).await.is_none());
@ -662,13 +525,13 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_queue_next_batch_min_size() { async fn test_queue_next_batch_min_size() {
let queue = default_queue(); let queue = Queue::new(false, 1, None, 0);
let (entry1, _guard1) = default_entry(); let (entry1, _guard1) = default_entry();
let (entry2, _guard2) = default_entry(); let (entry2, _guard2) = default_entry();
queue.append(entry1); queue.append(entry1);
queue.append(entry2); queue.append(entry2);
let (entries, batch, _) = queue.next_batch(None, None, 2, 4).await.unwrap(); let (entries, batch, _) = queue.next_batch(None, None, 2, 2).await.unwrap();
assert_eq!(entries.len(), 2); assert_eq!(entries.len(), 2);
assert!(entries.contains_key(&0)); assert!(entries.contains_key(&0));
assert!(entries.contains_key(&1)); assert!(entries.contains_key(&1));
@ -685,7 +548,7 @@ mod tests {
// Not enough token budget // Not enough token budget
assert!(queue.next_batch(Some(1), None, 0, 0).await.is_none()); assert!(queue.next_batch(Some(1), None, 0, 0).await.is_none());
// Ok // Ok
let (entries2, batch2, _) = queue.next_batch(Some(1), None, 2, 4).await.unwrap(); let (entries2, batch2, _) = queue.next_batch(Some(1), None, 2, 2).await.unwrap();
assert_eq!(entries2.len(), 1); assert_eq!(entries2.len(), 1);
assert!(entries2.contains_key(&2)); assert!(entries2.contains_key(&2));
assert!(entries2.get(&2).unwrap().batch_time.is_some()); assert!(entries2.get(&2).unwrap().batch_time.is_some());
@ -695,7 +558,7 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_queue_next_batch_max_size() { async fn test_queue_next_batch_max_size() {
let queue = default_queue(); let queue = Queue::new(false, 1, None, 0);
let (entry1, _guard1) = default_entry(); let (entry1, _guard1) = default_entry();
let (entry2, _guard2) = default_entry(); let (entry2, _guard2) = default_entry();
queue.append(entry1); queue.append(entry1);
@ -711,13 +574,13 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_queue_next_batch_token_budget() { async fn test_queue_next_batch_token_budget() {
let queue = default_queue(); let queue = Queue::new(false, 1, None, 0);
let (entry1, _guard1) = default_entry(); let (entry1, _guard1) = default_entry();
let (entry2, _guard2) = default_entry(); let (entry2, _guard2) = default_entry();
queue.append(entry1); queue.append(entry1);
queue.append(entry2); queue.append(entry2);
let (entries, batch, _) = queue.next_batch(None, None, 1, 2).await.unwrap(); let (entries, batch, _) = queue.next_batch(None, None, 1, 1).await.unwrap();
assert_eq!(entries.len(), 1); assert_eq!(entries.len(), 1);
assert!(entries.contains_key(&0)); assert!(entries.contains_key(&0));
assert_eq!(batch.id, 0); assert_eq!(batch.id, 0);
@ -726,7 +589,7 @@ mod tests {
let (entry3, _guard3) = default_entry(); let (entry3, _guard3) = default_entry();
queue.append(entry3); queue.append(entry3);
let (entries, batch, _) = queue.next_batch(None, None, 3, 6).await.unwrap(); let (entries, batch, _) = queue.next_batch(None, None, 3, 3).await.unwrap();
assert_eq!(entries.len(), 2); assert_eq!(entries.len(), 2);
assert!(entries.contains_key(&1)); assert!(entries.contains_key(&1));
assert!(entries.contains_key(&2)); assert!(entries.contains_key(&2));
@ -736,7 +599,7 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_queue_next_batch_token_speculate() { async fn test_queue_next_batch_token_speculate() {
let queue = Queue::new(true, 1, 2, 1, None, 2); let queue = Queue::new(false, 1, None, 2);
let (entry1, _guard1) = default_entry(); let (entry1, _guard1) = default_entry();
let (entry2, _guard2) = default_entry(); let (entry2, _guard2) = default_entry();
queue.append(entry1); queue.append(entry1);
@ -755,7 +618,7 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_queue_next_batch_dropped_receiver() { async fn test_queue_next_batch_dropped_receiver() {
let queue = default_queue(); let queue = Queue::new(false, 1, None, 0);
let (entry, _) = default_entry(); let (entry, _) = default_entry();
queue.append(entry); queue.append(entry);

4
router/src/server.rs
View File

@ -1,5 +1,3 @@
/// Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.
use crate::config::Config; use crate::config::Config;
/// HTTP Server logic /// HTTP Server logic
use crate::health::Health; use crate::health::Health;
@ -1493,8 +1491,6 @@ pub async fn run(
max_batch_size, max_batch_size,
max_concurrent_requests, max_concurrent_requests,
shard_info.requires_padding, shard_info.requires_padding,
max_input_length as u32,
max_total_tokens as u32,
shard_info.window_size, shard_info.window_size,
shard_info.speculate, shard_info.speculate,
generation_health, generation_health,

50
router/src/validation.rs
View File

@ -1,12 +1,9 @@
/// Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.
use crate::config::Config; use crate::config::Config;
/// Payload validation logic /// Payload validation logic
use crate::validation::ValidationError::{BestOfSampling, BestOfSeed, EmptyInput}; use crate::validation::ValidationError::{BestOfSampling, BestOfSeed, EmptyInput};
use crate::{GenerateParameters, GenerateRequest, GrammarType}; use crate::{GenerateParameters, GenerateRequest, GrammarType};
use jsonschema::{Draft, JSONSchema}; use jsonschema::{Draft, JSONSchema};
use rand::{thread_rng, Rng}; use rand::{thread_rng, Rng};
use std::{cmp, env};
use serde_json::Value; use serde_json::Value;
use std::io::Cursor; use std::io::Cursor;
use text_generation_client::{ use text_generation_client::{
@ -34,7 +31,6 @@ pub struct Validation {
disable_grammar_support: bool, disable_grammar_support: bool,
/// Channel to communicate with the background tokenization task /// Channel to communicate with the background tokenization task
sender: Option<mpsc::UnboundedSender<TokenizerRequest>>, sender: Option<mpsc::UnboundedSender<TokenizerRequest>>,
skip_tokenizer_in_tgi: bool,
} }
impl Validation { impl Validation {
@ -77,10 +73,6 @@ impl Validation {
None None
}; };
let skip_tokenizer_in_tgi = env::var("SKIP_TOKENIZER_IN_TGI")
.ok()
.map_or(false, |value| value.to_lowercase() == "true");
Self { Self {
max_best_of, max_best_of,
sender, sender,
@ -89,7 +81,6 @@ impl Validation {
max_input_length, max_input_length,
max_total_tokens, max_total_tokens,
disable_grammar_support, disable_grammar_support,
skip_tokenizer_in_tgi,
} }
} }
@ -128,13 +119,10 @@ impl Validation {
// If we have a fast tokenizer // If we have a fast tokenizer
if let Some((encoding, inputs)) = self.tokenize(inputs.clone(), truncate).await? { if let Some((encoding, inputs)) = self.tokenize(inputs.clone(), truncate).await? {
// Create response channel // Create response channel
let input_length = if self.skip_tokenizer_in_tgi { let input_length = if let Some(truncate) = truncate {
inputs.chars().filter(|&c| c == ',').count() + 1 std::cmp::min(encoding.len(), truncate)
} else { } else {
cmp::max( encoding.len()
encoding.len(),
truncate.unwrap_or(self.max_input_length)
)
}; };
// Get total tokens // Get total tokens
@ -177,18 +165,17 @@ impl Validation {
} else { } else {
return Err(ValidationError::UnsetMaxNewTokens); return Err(ValidationError::UnsetMaxNewTokens);
}; };
let input_length = if self.skip_tokenizer_in_tgi { let mut input_length = truncate.unwrap_or(self.max_input_length);
inputs.chars().filter(|&c| c == ',').count() + 1
} else {
truncate.unwrap_or(self.max_input_length)
};
// We don't have a tokenizer, therefore we have no idea how long is the query, let
// them through and hope for the best.
// Validate MaxNewTokens // Validate MaxNewTokens
if (input_length as u32 + max_new_tokens) > self.max_total_tokens as u32 { if (input_length as u32 + max_new_tokens) > self.max_total_tokens as u32 {
return Err(ValidationError::MaxNewTokens( input_length = input_length.saturating_sub(max_new_tokens as usize);
self.max_total_tokens - self.max_input_length, // return Err(ValidationError::MaxNewTokens(
max_new_tokens, // self.max_total_tokens - self.max_input_length,
)); // max_new_tokens,
// ));
} }
Ok((inputs, input_length, max_new_tokens)) Ok((inputs, input_length, max_new_tokens))
@ -248,14 +235,6 @@ impl Validation {
return Err(ValidationError::FrequencyPenalty); return Err(ValidationError::FrequencyPenalty);
} }
// TODO: enable watermark with fp8 quantization
let quantization_enabled = env::var("QUANT_CONFIG")
.ok()
.map_or(false, |value| !value.is_empty());
if watermark && quantization_enabled {
return Err(ValidationError::WatermarkWithQuantization);
}
// Different because the proto default value is not a valid value // Different because the proto default value is not a valid value
// for the user // for the user
let top_p = top_p let top_p = top_p
@ -730,8 +709,6 @@ pub enum ValidationError {
InvalidImageContent(String), InvalidImageContent(String),
#[error("Could not fetch image: {0}")] #[error("Could not fetch image: {0}")]
FailedFetchImage(#[from] reqwest::Error), FailedFetchImage(#[from] reqwest::Error),
#[error("`watermark` = true is not allowed with FP8 quantization.")]
WatermarkWithQuantization,
} }
#[cfg(test)] #[cfg(test)]
@ -768,7 +745,8 @@ mod tests {
.validate_input("Hello".to_string(), None, Some(max_new_tokens)) .validate_input("Hello".to_string(), None, Some(max_new_tokens))
.await .await
{ {
Err(ValidationError::MaxNewTokens(1, 10)) => (), // Err(ValidationError::MaxNewTokens(1, 10)) => (),
Ok((_s, 0, 10)) => (),
r => panic!("Unexpected not max new tokens: {r:?}"), r => panic!("Unexpected not max new tokens: {r:?}"),
} }
} }
@ -801,7 +779,7 @@ mod tests {
.validate_input("Hello".to_string(), None, Some(max_new_tokens)) .validate_input("Hello".to_string(), None, Some(max_new_tokens))
.await .await
{ {
Err(ValidationError::MaxTotalTokens(6, 5, 10)) => (), Err(ValidationError::MaxTotalTokens(6, 1, 10)) => (),
_ => panic!("Unexpected not max new tokens"), _ => panic!("Unexpected not max new tokens"),
} }
} }

14
server/Makefile
View File

@ -19,18 +19,12 @@ gen-server:
install: gen-server install: gen-server
pip install pip --upgrade pip install pip --upgrade
pip install -r requirements.txt pip install -r requirements_cuda.txt
pip install -e "." pip install -e ".[bnb, accelerate, quantize, peft, outlines]"
run-dev: run-dev:
SAFETENSORS_FAST_GPU=1 python -m torch.distributed.run --nproc_per_node=2 text_generation_server/cli.py serve bigscience/bloom-560m --sharded SAFETENSORS_FAST_GPU=1 python -m torch.distributed.run --nproc_per_node=2 text_generation_server/cli.py serve bigscience/bloom-560m --sharded
install-poetry:
curl -sSL https://install.python-poetry.org | python3 -
update-lock:
rm poetry.lock
poetry lock --no-update
export-requirements: export-requirements:
poetry export -o requirements.txt --without-hashes poetry export -o requirements_cuda.txt --without-hashes
poetry export -o requirements_rocm.txt --without-hashes

91
server/dill-0.3.7-patch.sh
View File

@ -1,91 +0,0 @@
#!/bin/bash
git clone -b dill-0.3.7 https://github.com/uqfoundation/dill.git
pushd dill
cat <<EOF > dill-0.3.7.patch
diff --git a/dill/_dill.py b/dill/_dill.py
index d0cf543..f6eb662 100644
--- a/dill/_dill.py
+++ b/dill/_dill.py
@@ -69,7 +69,15 @@ TypeType = type # 'new-style' classes #XXX: unregistered
XRangeType = range
from types import MappingProxyType as DictProxyType, new_class
from pickle import DEFAULT_PROTOCOL, HIGHEST_PROTOCOL, PickleError, PicklingError, UnpicklingError
-import __main__ as _main_module
+class _LazyMainModule(object):
+ _module = None
+ @property
+ def module(self):
+ if self._module is None:
+ import __main__ as _m_module
+ self._module = _m_module
+ return self._module
+_main_module = _LazyMainModule()
import marshal
import gc
# import zlib
@@ -353,7 +361,7 @@ class Pickler(StockPickler):
_fmode = kwds.pop('fmode', None)
_recurse = kwds.pop('recurse', None)
StockPickler.__init__(self, file, *args, **kwds)
- self._main = _main_module
+ self._main = _main_module.module
self._diff_cache = {}
self._byref = settings['byref'] if _byref is None else _byref
self._strictio = False #_strictio
@@ -435,12 +443,12 @@ class Unpickler(StockUnpickler):
settings = Pickler.settings
_ignore = kwds.pop('ignore', None)
StockUnpickler.__init__(self, *args, **kwds)
- self._main = _main_module
+ self._main = _main_module.module
self._ignore = settings['ignore'] if _ignore is None else _ignore
def load(self): #NOTE: if settings change, need to update attributes
obj = StockUnpickler.load(self)
- if type(obj).__module__ == getattr(_main_module, '__name__', '__main__'):
+ if type(obj).__module__ == getattr(self._main, '__name__', '__main__'):
if not self._ignore:
# point obj class to main
try: obj.__class__ = getattr(self._main, type(obj).__name__)
@@ -1194,11 +1202,11 @@ def save_module_dict(pickler, obj):
logger.trace(pickler, "D1: %s", _repr_dict(obj)) # obj
pickler.write(bytes('c__builtin__\n__main__\n', 'UTF-8'))
logger.trace(pickler, "# D1")
- elif (not is_dill(pickler, child=False)) and (obj == _main_module.__dict__):
+ elif (not is_dill(pickler, child=False)) and (obj == _main_module.module.__dict__):
logger.trace(pickler, "D3: %s", _repr_dict(obj)) # obj
pickler.write(bytes('c__main__\n__dict__\n', 'UTF-8')) #XXX: works in general?
logger.trace(pickler, "# D3")
- elif '__name__' in obj and obj != _main_module.__dict__ \\
+ elif '__name__' in obj and obj != _main_module.module.__dict__ \\
and type(obj['__name__']) is str \\
and obj is getattr(_import_module(obj['__name__'],True), '__dict__', None):
logger.trace(pickler, "D4: %s", _repr_dict(obj)) # obj
diff --git a/dill/session.py b/dill/session.py
index 74234ab..1be8d89 100644
--- a/dill/session.py
+++ b/dill/session.py
@@ -233,7 +233,7 @@ def dump_module(
protocol = settings['protocol']
main = module
if main is None:
- main = _main_module
+ main = _main_module.module
elif isinstance(main, str):
main = _import_module(main)
if not isinstance(main, ModuleType):
@@ -501,7 +501,7 @@ def load_module(
pass
assert loaded is main
_restore_modules(unpickler, main)
- if main is _main_module or main is module:
+ if main is _main_module.module or main is module:
return None
else:
return main
EOF
git apply dill-0.3.7.patch
python -m pip install .
popd
rm -fr dill

91
server/dill-0.3.8-patch.sh
View File

@ -1,91 +0,0 @@
#!/bin/bash
git clone -b 0.3.8 https://github.com/uqfoundation/dill.git
pushd dill
cat <<EOF > dill-0.3.8.patch
diff --git a/dill/_dill.py b/dill/_dill.py
index d42432f..1d251e6 100644
--- a/dill/_dill.py
+++ b/dill/_dill.py
@@ -69,7 +69,15 @@ TypeType = type # 'new-style' classes #XXX: unregistered
XRangeType = range
from types import MappingProxyType as DictProxyType, new_class
from pickle import DEFAULT_PROTOCOL, HIGHEST_PROTOCOL, PickleError, PicklingError, UnpicklingError
-import __main__ as _main_module
+class _LazyMainModule(object):
+ _module = None
+ @property
+ def module(self):
+ if self._module is None:
+ import __main__ as _m_module
+ self._module = _m_module
+ return self._module
+_main_module = _LazyMainModule()
import marshal
import gc
# import zlib
@@ -355,7 +363,7 @@ class Pickler(StockPickler):
_fmode = kwds.pop('fmode', None)
_recurse = kwds.pop('recurse', None)
StockPickler.__init__(self, file, *args, **kwds)
- self._main = _main_module
+ self._main = _main_module.module
self._diff_cache = {}
self._byref = settings['byref'] if _byref is None else _byref
self._strictio = False #_strictio
@@ -437,12 +445,12 @@ class Unpickler(StockUnpickler):
settings = Pickler.settings
_ignore = kwds.pop('ignore', None)
StockUnpickler.__init__(self, *args, **kwds)
- self._main = _main_module
+ self._main = _main_module.module
self._ignore = settings['ignore'] if _ignore is None else _ignore
def load(self): #NOTE: if settings change, need to update attributes
obj = StockUnpickler.load(self)
- if type(obj).__module__ == getattr(_main_module, '__name__', '__main__'):
+ if type(obj).__module__ == getattr(self._main, '__name__', '__main__'):
if not self._ignore:
# point obj class to main
try: obj.__class__ = getattr(self._main, type(obj).__name__)
@@ -1199,11 +1207,11 @@ def save_module_dict(pickler, obj):
logger.trace(pickler, "D1: %s", _repr_dict(obj)) # obj
pickler.write(bytes('c__builtin__\n__main__\n', 'UTF-8'))
logger.trace(pickler, "# D1")
- elif (not is_dill(pickler, child=False)) and (obj == _main_module.__dict__):
+ elif (not is_dill(pickler, child=False)) and (obj == _main_module.module.__dict__):
logger.trace(pickler, "D3: %s", _repr_dict(obj)) # obj
pickler.write(bytes('c__main__\n__dict__\n', 'UTF-8')) #XXX: works in general?
logger.trace(pickler, "# D3")
- elif '__name__' in obj and obj != _main_module.__dict__ \\
+ elif '__name__' in obj and obj != _main_module.module.__dict__ \\
and type(obj['__name__']) is str \\
and obj is getattr(_import_module(obj['__name__'],True), '__dict__', None):
logger.trace(pickler, "D4: %s", _repr_dict(obj)) # obj
diff --git a/dill/session.py b/dill/session.py
index e91068a..a921b43 100644
--- a/dill/session.py
+++ b/dill/session.py
@@ -233,7 +233,7 @@ def dump_module(
protocol = settings['protocol']
main = module
if main is None:
- main = _main_module
+ main = _main_module.module
elif isinstance(main, str):
main = _import_module(main)
if not isinstance(main, ModuleType):
@@ -501,7 +501,7 @@ def load_module(
pass
assert loaded is main
_restore_modules(unpickler, main)
- if main is _main_module or main is module:
+ if main is _main_module.module or main is module:
return None
else:
return main
EOF
git apply dill-0.3.8.patch
python -m pip install .
popd
rm -fr dill

2821
server/poetry.lock generated
View File

File diff suppressed because it is too large Load Diff

46
server/pyproject.toml
View File

@ -9,34 +9,58 @@ text-generation-server = 'text_generation_server.cli:app'
[tool.poetry.dependencies] [tool.poetry.dependencies]
python = ">=3.9,<3.13" python = ">=3.9,<3.13"
protobuf = "^3.20.3" protobuf = "^4.21.7"
grpcio = "^1.51.1" grpcio = "^1.51.1"
grpcio-status = "*" grpcio-status = "^1.51.1"
grpcio-reflection = "*" grpcio-reflection = "^1.51.1"
grpc-interceptor = "^0.15.0" grpc-interceptor = "^0.15.0"
typer = "^0.7.0" typer = "^0.6.1"
accelerate = { version = "^0.29.1", optional = true }
bitsandbytes = { version = "^0.43.0", optional = true }
safetensors = "^0.4"
loguru = "^0.6.0" loguru = "^0.6.0"
opentelemetry-api = "^1.15.0" opentelemetry-api = "^1.15.0"
opentelemetry-exporter-otlp = "^1.15.0" opentelemetry-exporter-otlp = "^1.15.0"
opentelemetry-instrumentation-grpc = "^0.36b0" opentelemetry-instrumentation-grpc = "^0.36b0"
hf-transfer = "^0.1.2" hf-transfer = "^0.1.2"
sentencepiece = "^0.1.97" sentencepiece = "^0.1.97"
peft = "^0.10" tokenizers = "^0.19.1"
optimum-habana = "1.13.2" huggingface-hub = "^0.23"
transformers = "4.43.4" transformers = "^4.41"
numpy = "1.26.4" einops = "^0.6.1"
accelerate = "0.33.0" texttable = { version = "^1.6.7", optional = true }
datasets = { version = "^2.14.0", optional = true }
peft = { version = "^0.10", optional = true }
torch = { version = "^2.3.0", optional = true }
scipy = "^1.11.1"
pillow = "^10.0.0"
outlines= { version = "^0.0.36", optional = true } outlines= { version = "^0.0.36", optional = true }
prometheus-client = "^0.20.0" prometheus-client = "^0.20.0"
py-cpuinfo = "^9.0.0" py-cpuinfo = "^9.0.0"
[tool.poetry.extras]
torch = ["torch"]
accelerate = ["accelerate"]
bnb = ["bitsandbytes"]
peft = ["peft"]
quantize = ["texttable", "datasets", "accelerate"]
outlines = ["outlines"]
[tool.poetry.group.dev.dependencies] [tool.poetry.group.dev.dependencies]
grpcio-tools = "*" grpcio-tools = "^1.51.1"
pytest = "^7.3.0" pytest = "^7.3.0"
[[tool.poetry.source]]
name = "pytorch-gpu-src"
url = "https://download.pytorch.org/whl/cu121"
priority = "explicit"
[tool.pytest.ini_options] [tool.pytest.ini_options]
markers = ["private: marks tests as requiring an admin hf token (deselect with '-m \"not private\"')"] markers = ["private: marks tests as requiring an admin hf token (deselect with '-m \"not private\"')"]
[build-system] [build-system]
requires = ["poetry-core>=1.0.0"] requires = [
"poetry-core>=1.0.0",
]
build-backend = "poetry.core.masonry.api" build-backend = "poetry.core.masonry.api"

88
server/requirements.txt
View File

@ -1,88 +0,0 @@
accelerate==0.33.0 ; python_version >= "3.9" and python_version < "3.13"
aiohappyeyeballs==2.4.0 ; python_version >= "3.9" and python_version < "3.13"
aiohttp==3.10.5 ; python_version >= "3.9" and python_version < "3.13"
aiosignal==1.3.1 ; python_version >= "3.9" and python_version < "3.13"
async-timeout==4.0.3 ; python_version >= "3.9" and python_version < "3.11"
attrs==24.2.0 ; python_version >= "3.9" and python_version < "3.13"
backoff==2.2.1 ; python_version >= "3.9" and python_version < "3.13"
certifi==2024.7.4 ; python_version >= "3.9" and python_version < "3.13"
charset-normalizer==3.3.2 ; python_version >= "3.9" and python_version < "3.13"
click==8.1.7 ; python_version >= "3.9" and python_version < "3.13"
colorama==0.4.6 ; python_version >= "3.9" and python_version < "3.13" and (sys_platform == "win32" or platform_system == "Windows")
coloredlogs==15.0.1 ; python_version >= "3.9" and python_version < "3.13"
datasets==2.21.0 ; python_version >= "3.9" and python_version < "3.13"
deprecated==1.2.14 ; python_version >= "3.9" and python_version < "3.13"
diffusers==0.29.2 ; python_version >= "3.9" and python_version < "3.13"
dill==0.3.8 ; python_version >= "3.9" and python_version < "3.13"
filelock==3.15.4 ; python_version >= "3.9" and python_version < "3.13"
frozenlist==1.4.1 ; python_version >= "3.9" and python_version < "3.13"
fsspec==2024.6.1 ; python_version >= "3.9" and python_version < "3.13"
fsspec[http]==2024.6.1 ; python_version >= "3.9" and python_version < "3.13"
googleapis-common-protos==1.63.2 ; python_version >= "3.9" and python_version < "3.13"
grpc-interceptor==0.15.4 ; python_version >= "3.9" and python_version < "3.13"
grpcio-reflection==1.48.2 ; python_version >= "3.9" and python_version < "3.13"
grpcio-status==1.48.2 ; python_version >= "3.9" and python_version < "3.13"
grpcio==1.66.0 ; python_version >= "3.9" and python_version < "3.13"
hf-transfer==0.1.8 ; python_version >= "3.9" and python_version < "3.13"
huggingface-hub==0.24.6 ; python_version >= "3.9" and python_version < "3.13"
humanfriendly==10.0 ; python_version >= "3.9" and python_version < "3.13"
idna==3.8 ; python_version >= "3.9" and python_version < "3.13"
importlib-metadata==8.4.0 ; python_version >= "3.9" and python_version < "3.13"
jinja2==3.1.4 ; python_version >= "3.9" and python_version < "3.13"
joblib==1.4.2 ; python_version >= "3.9" and python_version < "3.13"
loguru==0.6.0 ; python_version >= "3.9" and python_version < "3.13"
markupsafe==2.1.5 ; python_version >= "3.9" and python_version < "3.13"
mpmath==1.3.0 ; python_version >= "3.9" and python_version < "3.13"
multidict==6.0.5 ; python_version >= "3.9" and python_version < "3.13"
multiprocess==0.70.16 ; python_version >= "3.9" and python_version < "3.13"
networkx==3.2.1 ; python_version >= "3.9" and python_version < "3.13"
numpy==1.26.4 ; python_version >= "3.9" and python_version < "3.13"
opentelemetry-api==1.15.0 ; python_version >= "3.9" and python_version < "3.13"
opentelemetry-exporter-otlp-proto-grpc==1.15.0 ; python_version >= "3.9" and python_version < "3.13"
opentelemetry-exporter-otlp-proto-http==1.15.0 ; python_version >= "3.9" and python_version < "3.13"
opentelemetry-exporter-otlp==1.15.0 ; python_version >= "3.9" and python_version < "3.13"
opentelemetry-instrumentation-grpc==0.36b0 ; python_version >= "3.9" and python_version < "3.13"
opentelemetry-instrumentation==0.36b0 ; python_version >= "3.9" and python_version < "3.13"
opentelemetry-proto==1.15.0 ; python_version >= "3.9" and python_version < "3.13"
opentelemetry-sdk==1.15.0 ; python_version >= "3.9" and python_version < "3.13"
opentelemetry-semantic-conventions==0.36b0 ; python_version >= "3.9" and python_version < "3.13"
optimum-habana==1.13.2 ; python_version >= "3.9" and python_version < "3.13"
optimum==1.21.4 ; python_version >= "3.9" and python_version < "3.13"
packaging==24.1 ; python_version >= "3.9" and python_version < "3.13"
pandas==2.2.2 ; python_version >= "3.9" and python_version < "3.13"
peft==0.10.0 ; python_version >= "3.9" and python_version < "3.13"
pillow==10.4.0 ; python_version >= "3.9" and python_version < "3.13"
prometheus-client==0.20.0 ; python_version >= "3.9" and python_version < "3.13"
protobuf==3.20.3 ; python_version >= "3.9" and python_version < "3.13"
psutil==6.0.0 ; python_version >= "3.9" and python_version < "3.13"
py-cpuinfo==9.0.0 ; python_version >= "3.9" and python_version < "3.13"
pyarrow==17.0.0 ; python_version >= "3.9" and python_version < "3.13"
pyreadline3==3.4.1 ; sys_platform == "win32" and python_version >= "3.9" and python_version < "3.13"
python-dateutil==2.9.0.post0 ; python_version >= "3.9" and python_version < "3.13"
pytz==2024.1 ; python_version >= "3.9" and python_version < "3.13"
pyyaml==6.0.2 ; python_version >= "3.9" and python_version < "3.13"
regex==2024.7.24 ; python_version >= "3.9" and python_version < "3.13"
requests==2.32.3 ; python_version >= "3.9" and python_version < "3.13"
safetensors==0.4.4 ; python_version >= "3.9" and python_version < "3.13"
scikit-learn==1.5.1 ; python_version >= "3.9" and python_version < "3.13"
scipy==1.13.1 ; python_version >= "3.9" and python_version < "3.13"
sentence-transformers[train]==3.0.1 ; python_version >= "3.9" and python_version < "3.13"
sentencepiece==0.1.99 ; python_version >= "3.9" and python_version < "3.13"
setuptools==73.0.1 ; python_version >= "3.9" and python_version < "3.13"
six==1.16.0 ; python_version >= "3.9" and python_version < "3.13"
sympy==1.12.1 ; python_version >= "3.9" and python_version < "3.13"
threadpoolctl==3.5.0 ; python_version >= "3.9" and python_version < "3.13"
tokenizers==0.19.1 ; python_version >= "3.9" and python_version < "3.13"
tqdm==4.66.5 ; python_version >= "3.9" and python_version < "3.13"
transformers==4.43.4 ; python_version >= "3.9" and python_version < "3.13"
transformers[sentencepiece]==4.43.4 ; python_version >= "3.9" and python_version < "3.13"
triton==3.0.0 ; platform_system == "Linux" and platform_machine == "x86_64" and python_version < "3.13" and python_version >= "3.9"
typer==0.7.0 ; python_version >= "3.9" and python_version < "3.13"
typing-extensions==4.12.2 ; python_version >= "3.9" and python_version < "3.13"
tzdata==2024.1 ; python_version >= "3.9" and python_version < "3.13"
urllib3==2.2.2 ; python_version >= "3.9" and python_version < "3.13"
win32-setctime==1.1.0 ; python_version >= "3.9" and python_version < "3.13" and sys_platform == "win32"
wrapt==1.16.0 ; python_version >= "3.9" and python_version < "3.13"
xxhash==3.5.0 ; python_version >= "3.9" and python_version < "3.13"
yarl==1.9.4 ; python_version >= "3.9" and python_version < "3.13"
zipp==3.20.0 ; python_version >= "3.9" and python_version < "3.13"

43
server/tests/models/test_bloom.py
View File

@ -1,5 +1,3 @@
# Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.
import pytest import pytest
import torch import torch
@ -7,9 +5,9 @@ from copy import copy
from transformers import AutoTokenizer from transformers import AutoTokenizer
from text_generation_server.pb import generate_pb2 from text_generation_server.pb import generate_pb2
from text_generation_server.models.causal_lm import CausalLMBatch, PAD_SEQUENCE_TO_MULTIPLE_OF from text_generation_server.models.causal_lm import CausalLMBatch
from text_generation_server.utils import weight_hub_files, download_weights from text_generation_server.utils import weight_hub_files, download_weights
from text_generation_server.models.bloom import BloomCausalLMBatch, BLOOM from text_generation_server.models.bloom import BloomCausalLMBatch, BLOOMSharded
@pytest.fixture(scope="session") @pytest.fixture(scope="session")
@ -18,7 +16,7 @@ def default_bloom():
revision = "main" revision = "main"
filenames = weight_hub_files(model_id, revision, ".safetensors") filenames = weight_hub_files(model_id, revision, ".safetensors")
download_weights(filenames, model_id, revision) download_weights(filenames, model_id, revision)
return BLOOM(model_id) return BLOOMSharded(model_id)
@pytest.fixture(scope="session") @pytest.fixture(scope="session")
@ -32,7 +30,7 @@ def default_pb_request(default_pb_parameters, default_pb_stop_parameters):
id=0, id=0,
inputs="Test", inputs="Test",
prefill_logprobs=True, prefill_logprobs=True,
truncate=PAD_SEQUENCE_TO_MULTIPLE_OF, truncate=100,
parameters=default_pb_parameters, parameters=default_pb_parameters,
stopping_parameters=default_pb_stop_parameters, stopping_parameters=default_pb_stop_parameters,
) )
@ -46,7 +44,7 @@ def default_pb_batch(default_pb_request):
@pytest.fixture @pytest.fixture
def default_bloom_batch(default_pb_batch, bloom_560m_tokenizer): def default_bloom_batch(default_pb_batch, bloom_560m_tokenizer):
return BloomCausalLMBatch.from_pb( return BloomCausalLMBatch.from_pb(
default_pb_batch, bloom_560m_tokenizer, torch.float32, torch.device("hpu") default_pb_batch, bloom_560m_tokenizer, torch.float32, torch.device("cpu")
) )
@ -60,7 +58,7 @@ def default_multi_requests_bloom_batch(default_pb_request, bloom_560m_tokenizer)
batch_pb = generate_pb2.Batch(id=0, requests=[req_0, req_1], size=2) batch_pb = generate_pb2.Batch(id=0, requests=[req_0, req_1], size=2)
return BloomCausalLMBatch.from_pb( return BloomCausalLMBatch.from_pb(
batch_pb, bloom_560m_tokenizer, torch.float32, torch.device("hpu") batch_pb, bloom_560m_tokenizer, torch.float32, torch.device("cpu")
) )
@ -68,29 +66,30 @@ def test_batch_from_pb(default_pb_batch, default_bloom_batch):
batch = default_bloom_batch batch = default_bloom_batch
assert batch.batch_id == default_pb_batch.id assert batch.batch_id == default_pb_batch.id
assert len(batch.requests) == len(default_pb_batch.requests) == default_pb_batch.size assert batch.requests == default_pb_batch.requests
for request, pb_request in zip(batch.requests, default_pb_batch.requests):
assert request.data == pb_request
assert batch.input_ids[0][-1] == 3 assert len(batch.input_ids) == default_pb_batch.size
assert batch.input_ids[0][-2] == 10264 assert batch.input_ids[0][-1] == 10264
assert torch.all(batch.input_ids[0][:-2] == 3) assert torch.all(batch.input_ids[0][:-1] == 3)
assert batch.attention_mask[0][-1] == 0 assert batch.attention_mask[0][0] == 1
assert batch.attention_mask[0][-2] == 1 assert torch.all(batch.attention_mask[0][1:] == 0)
assert torch.all(batch.attention_mask[0][:-2] == 0)
assert batch.past_key_values is None assert batch.past_key_values is None
assert all( assert all(
[ [
torch.equal(input_ids, request.all_input_ids[:batch.input_length+1, 0]) torch.equal(input_ids, all_input_ids[:, 0])
for input_ids, request in zip(batch.input_ids, batch.requests) for input_ids, all_input_ids in zip(batch.input_ids, batch.all_input_ids)
] ]
) )
assert batch.input_lengths == [1]
assert len(batch) == default_pb_batch.size assert len(batch) == default_pb_batch.size
assert batch.max_input_length == batch.input_length == PAD_SEQUENCE_TO_MULTIPLE_OF - 1 assert len(batch.next_token_choosers) == len(batch.stopping_criterias) == len(batch)
assert batch.max_input_length == batch.input_lengths[0]
def test_batch_concatenate_no_prefill(default_bloom_batch): def test_batch_concatenate_no_prefill(default_bloom_batch):
@ -102,7 +101,6 @@ def test_causal_lm_batch_type(default_bloom):
assert default_bloom.batch_type == BloomCausalLMBatch assert default_bloom.batch_type == BloomCausalLMBatch
@pytest.mark.skip
def test_causal_lm_generate_token(default_bloom, default_bloom_batch): def test_causal_lm_generate_token(default_bloom, default_bloom_batch):
sequence_length = len(default_bloom_batch.all_input_ids[0]) sequence_length = len(default_bloom_batch.all_input_ids[0])
generations, next_batch, _ = default_bloom.generate_token(default_bloom_batch) generations, next_batch, _ = default_bloom.generate_token(default_bloom_batch)
@ -152,7 +150,6 @@ def test_causal_lm_generate_token(default_bloom, default_bloom_batch):
assert generations[0].request_id == 0 assert generations[0].request_id == 0
@pytest.mark.skip
def test_causal_lm_generate_token_completion(default_bloom, default_bloom_batch): def test_causal_lm_generate_token_completion(default_bloom, default_bloom_batch):
next_batch = default_bloom_batch next_batch = default_bloom_batch
for _ in range(default_bloom_batch.stopping_criterias[0].max_new_tokens - 1): for _ in range(default_bloom_batch.stopping_criterias[0].max_new_tokens - 1):
@ -173,7 +170,6 @@ def test_causal_lm_generate_token_completion(default_bloom, default_bloom_batch)
) )
@pytest.mark.skip
def test_causal_lm_generate_token_completion_multi( def test_causal_lm_generate_token_completion_multi(
default_bloom, default_multi_requests_bloom_batch default_bloom, default_multi_requests_bloom_batch
): ):
@ -224,7 +220,6 @@ def test_causal_lm_generate_token_completion_multi(
) )
@pytest.mark.skip
def test_batch_concatenate( def test_batch_concatenate(
default_bloom, default_bloom_batch, default_multi_requests_bloom_batch default_bloom, default_bloom_batch, default_multi_requests_bloom_batch
): ):

290
server/tests/models/test_causal_lm.py
View File

@ -1,5 +1,3 @@
# Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.
import pytest import pytest
import torch import torch
@ -7,27 +5,19 @@ from copy import copy
from transformers import AutoTokenizer from transformers import AutoTokenizer
from text_generation_server.pb import generate_pb2 from text_generation_server.pb import generate_pb2
from text_generation_server.models import get_model from text_generation_server.models.causal_lm import CausalLM, CausalLMBatch
from text_generation_server.models.causal_lm import (
CausalLMBatch,
PREFILL_BATCH_BUCKET_SIZE,
PAD_SEQUENCE_TO_MULTIPLE_OF,
MAX_TOTAL_TOKENS,
BATCH_BUCKET_SIZE,
)
PAD_TOKEN=0
@pytest.fixture(scope="session") @pytest.fixture(scope="session")
def default_causal_lm(): def default_causal_lm():
return get_model("meta-llama/Llama-2-7b-hf", None, None, None, None) return CausalLM("gpt2")
@pytest.fixture(scope="session") @pytest.fixture(scope="session")
def default_tokenizer(default_causal_lm): def gpt2_tokenizer():
default_causal_lm.tokenizer.pad_token_id = PAD_TOKEN tokenizer = AutoTokenizer.from_pretrained("gpt2", padding_side="left")
return default_causal_lm.tokenizer tokenizer.pad_token_id = 50256
return tokenizer
@pytest.fixture @pytest.fixture
@ -36,7 +26,7 @@ def default_pb_request(default_pb_parameters, default_pb_stop_parameters):
id=0, id=0,
inputs="Test", inputs="Test",
prefill_logprobs=True, prefill_logprobs=True,
truncate=PAD_SEQUENCE_TO_MULTIPLE_OF, truncate=100,
parameters=default_pb_parameters, parameters=default_pb_parameters,
stopping_parameters=default_pb_stop_parameters, stopping_parameters=default_pb_stop_parameters,
) )
@ -48,14 +38,14 @@ def default_pb_batch(default_pb_request):
@pytest.fixture @pytest.fixture
def default_causal_lm_batch(default_pb_batch, default_tokenizer): def default_causal_lm_batch(default_pb_batch, gpt2_tokenizer):
return CausalLMBatch.from_pb( return CausalLMBatch.from_pb(
default_pb_batch, default_tokenizer, torch.float32, torch.device("hpu") default_pb_batch, gpt2_tokenizer, torch.float32, torch.device("cpu")
) )
@pytest.fixture @pytest.fixture
def default_multi_requests_causal_lm_batch(default_pb_request, default_tokenizer): def default_multi_requests_causal_lm_batch(default_pb_request, gpt2_tokenizer):
req_0 = copy(default_pb_request) req_0 = copy(default_pb_request)
req_0.id = 1 req_0.id = 1
req_1 = default_pb_request req_1 = default_pb_request
@ -64,7 +54,7 @@ def default_multi_requests_causal_lm_batch(default_pb_request, default_tokenizer
batch_pb = generate_pb2.Batch(id=1, requests=[req_0, req_1], size=2) batch_pb = generate_pb2.Batch(id=1, requests=[req_0, req_1], size=2)
return CausalLMBatch.from_pb( return CausalLMBatch.from_pb(
batch_pb, default_tokenizer, torch.float32, torch.device("hpu") batch_pb, gpt2_tokenizer, torch.float32, torch.device("cpu")
) )
@ -72,37 +62,30 @@ def test_batch_from_pb(default_pb_batch, default_causal_lm_batch):
batch = default_causal_lm_batch batch = default_causal_lm_batch
assert batch.batch_id == default_pb_batch.id assert batch.batch_id == default_pb_batch.id
assert len(batch.requests) == len(default_pb_batch.requests) assert batch.requests == default_pb_batch.requests
for r in range(0,len(default_pb_batch.requests)): assert len(batch.input_ids) == default_pb_batch.size
assert batch.requests[r].data == default_pb_batch.requests[r] assert batch.input_ids[0][-1] == 14402
assert torch.all(batch.input_ids[0][:-1] == 50256)
# For Gaudi we are adding padding of multiplication of bucket size assert batch.attention_mask[0, 0] == 1
size_of_padded_to_bucket = ((default_pb_batch.size + PREFILL_BATCH_BUCKET_SIZE - 1) // PREFILL_BATCH_BUCKET_SIZE) * PREFILL_BATCH_BUCKET_SIZE assert torch.all(batch.attention_mask[0, 1:] == 0)
assert len(batch.input_ids) == size_of_padded_to_bucket
assert batch.input_ids[0][-2] == 4321
assert batch.input_ids[0][-3] == 1
assert torch.all(batch.input_ids[0][:-3] == PAD_TOKEN)
assert batch.input_ids[0][-1] == PAD_TOKEN
assert batch.attention_mask[0][-1] == 0
assert batch.attention_mask[0, -2] == 1
assert batch.attention_mask[0, -3] == 1
assert torch.all(batch.attention_mask[0, :-3] == 0)
assert batch.past_key_values is None assert batch.past_key_values is None
assert all( assert all(
[ [
torch.equal(input_ids.to('cpu'), request.all_input_ids[:batch.input_length + 1, 0]) torch.equal(input_ids, all_input_ids[:, 0])
for input_ids, request in zip(batch.input_ids, batch.requests) for input_ids, all_input_ids in zip(batch.input_ids, batch.all_input_ids)
] ]
) )
assert len(batch) == default_pb_batch.size assert batch.input_lengths == [1]
assert batch.max_input_length + 1 == default_pb_batch.requests[0].truncate assert len(batch) == default_pb_batch.size
assert len(batch.next_token_choosers) == len(batch.stopping_criterias) == len(batch)
assert batch.max_input_length == batch.input_lengths[0]
def test_batch_concatenate_no_prefill(default_causal_lm_batch): def test_batch_concatenate_no_prefill(default_causal_lm_batch):
@ -115,66 +98,73 @@ def test_causal_lm_batch_type(default_causal_lm):
def test_causal_lm_generate_token(default_causal_lm, default_causal_lm_batch): def test_causal_lm_generate_token(default_causal_lm, default_causal_lm_batch):
sequence_length = len(default_causal_lm_batch.all_input_ids[0])
generations, next_batch, _ = default_causal_lm.generate_token(
default_causal_lm_batch
)
sequence_length = len(default_causal_lm_batch.requests[0].all_input_ids) assert len(generations) == len(next_batch)
generations, next_batch, _ = default_causal_lm.generate_token([default_causal_lm_batch])
padding = next_batch.requests[0].stopping_criteria.max_new_tokens
assert isinstance(next_batch, CausalLMBatch) assert isinstance(next_batch, CausalLMBatch)
assert len(next_batch.attention_mask[0]) == PAD_SEQUENCE_TO_MULTIPLE_OF
assert next_batch.requests[0].all_input_ids[-padding-2] == 4321
assert torch.all(next_batch.requests[0].all_input_ids[-padding-1:] == PAD_TOKEN) assert len(next_batch.all_input_ids) == len(next_batch)
assert torch.all(next_batch.requests[0].all_input_ids[:-padding-3] == PAD_TOKEN) assert len(next_batch.all_input_ids[0]) == sequence_length + 1
assert len(next_batch.attention_mask[0]) == 11
assert next_batch.all_input_ids[0][-1] == 13
assert next_batch.all_input_ids[0][-2] == 14402
assert torch.all(next_batch.all_input_ids[0][:-2] == 50256)
generations, next_batch, _ = default_causal_lm.generate_token([default_causal_lm_batch]) assert torch.all(next_batch.attention_mask[0][0:2] == 1)
assert torch.all(next_batch.attention_mask[0][PAD_SEQUENCE_TO_MULTIPLE_OF-3:PAD_SEQUENCE_TO_MULTIPLE_OF] == 1) assert torch.all(next_batch.attention_mask[0][2:] == 0)
assert torch.all(next_batch.attention_mask[0][:PAD_SEQUENCE_TO_MULTIPLE_OF-3] == 0)
assert torch.all(next_batch.attention_mask[0][PAD_SEQUENCE_TO_MULTIPLE_OF+1:] == 0)
assert next_batch.requests[0].all_input_ids[-padding-2] == 4321 assert next_batch.input_ids.shape == (len(next_batch), 1)
assert next_batch.requests[0].all_input_ids[-padding-1] == 292 assert next_batch.input_ids[0, 0] == 13
assert torch.all(next_batch.requests[0].all_input_ids[-padding:] == PAD_TOKEN)
assert torch.all(next_batch.requests[0].all_input_ids[:-padding-3] == PAD_TOKEN)
assert next_batch.input_length == PAD_SEQUENCE_TO_MULTIPLE_OF assert next_batch.input_lengths == [2]
assert next_batch.max_input_length == next_batch.input_length assert next_batch.max_input_length == next_batch.input_lengths[0]
assert next_batch.past_key_values is not None assert next_batch.past_key_values is not None
assert all( assert all(
[p[0].shape == (BATCH_BUCKET_SIZE, 32, MAX_TOTAL_TOKENS, 128) for p in next_batch.past_key_values] [p[0].shape == (1, 12, sequence_length, 64) for p in next_batch.past_key_values]
) )
assert all( assert all(
[p[1].shape == (BATCH_BUCKET_SIZE, 32, MAX_TOTAL_TOKENS, 128) for p in next_batch.past_key_values] [p[1].shape == (1, 12, sequence_length, 64) for p in next_batch.past_key_values]
) )
assert all([generation.generated_text is None for generation in generations]) assert all([generation.generated_text is None for generation in generations])
assert all([len(generation.prefill_tokens) == PAD_SEQUENCE_TO_MULTIPLE_OF-1 for generation in generations]) assert all([len(generation.prefill_tokens) == 1 for generation in generations])
assert all([generation.tokens.token_ids[0] == 292 for generation in generations]) assert all(
assert all([generation.tokens.texts[0] == "ing" for generation in generations]) [
token_id.item() == 13
for generation in generations
for token_id in generation.tokens.token_ids
]
)
assert all(
[
token_text == "."
for generation in generations
for token_text in generation.tokens.texts
]
)
assert generations[0].request_id == 0 assert generations[0].request_id == 0
def test_causal_lm_generate_token_completion( def test_causal_lm_generate_token_completion(
default_causal_lm, default_causal_lm_batch default_causal_lm, default_causal_lm_batch
): ):
next_batch = default_causal_lm_batch next_batch = default_causal_lm_batch
generations, next_batch, _ = default_causal_lm.generate_token([next_batch]) for _ in range(default_causal_lm_batch.stopping_criterias[0].max_new_tokens - 1):
generations, next_batch, _ = default_causal_lm.generate_token(next_batch)
for _ in range(default_causal_lm_batch.requests[0].stopping_criteria.max_new_tokens - 1):
generations, next_batch, _ = default_causal_lm.generate_token([next_batch])
assert len(generations) == len(next_batch) assert len(generations) == len(next_batch)
generations, next_batch, _ = default_causal_lm.generate_token([next_batch]) generations, next_batch, _ = default_causal_lm.generate_token(next_batch)
assert next_batch is None assert next_batch is None
assert len(generations) == 1 assert len(generations) == 1
assert generations[0].generated_text.text == "ing the effect of a new method for the detection" assert generations[0].generated_text.text == ".java:784) at net.minecraft."
assert generations[0].request_id == default_causal_lm_batch.requests[0].data.id assert generations[0].request_id == default_causal_lm_batch.requests[0].id
assert ( assert (
generations[0].generated_text.generated_tokens generations[0].generated_text.generated_tokens
== default_causal_lm_batch.requests[0].stopping_criteria.max_new_tokens == default_causal_lm_batch.stopping_criterias[0].max_new_tokens
) )
@ -182,49 +172,51 @@ def test_causal_lm_generate_token_completion_multi(
default_causal_lm, default_multi_requests_causal_lm_batch default_causal_lm, default_multi_requests_causal_lm_batch
): ):
next_batch = default_multi_requests_causal_lm_batch next_batch = default_multi_requests_causal_lm_batch
generations, next_batch, _ = default_causal_lm.generate_token([next_batch])
for i in range( for i in range(
default_multi_requests_causal_lm_batch.requests[1].stopping_criteria.max_new_tokens - 1 default_multi_requests_causal_lm_batch.stopping_criterias[1].max_new_tokens - 1
): ):
generations, next_batch, _ = default_causal_lm.generate_token([next_batch]) generations, next_batch, _ = default_causal_lm.generate_token(next_batch)
assert len(generations) == len(next_batch) assert len(generations) == len(next_batch)
generations, next_batch, _ = default_causal_lm.generate_token([next_batch]) generations, next_batch, _ = default_causal_lm.generate_token(next_batch)
assert next_batch is not None assert next_batch is not None
assert len(generations) == 2 assert len(generations) == 2
assert generations[1].generated_text.text == "ing the effect of a" assert generations[1].generated_text.text == ".java:784)"
assert ( assert (
generations[1].request_id generations[1].request_id
== default_multi_requests_causal_lm_batch.requests[1].data.id == default_multi_requests_causal_lm_batch.requests[1].id
) )
assert ( assert (
generations[1].generated_text.generated_tokens generations[1].generated_text.generated_tokens
== default_multi_requests_causal_lm_batch.requests[1].stopping_criteria.max_new_tokens == default_multi_requests_causal_lm_batch.stopping_criterias[1].max_new_tokens
)
# Copy stopping_criterias before filtering
stopping_criterias = (
default_multi_requests_causal_lm_batch.stopping_criterias.copy()
) )
next_batch = next_batch.filter([next_batch.requests[0].data.id]) next_batch = next_batch.filter([next_batch.requests[0].id])
for _ in range( for _ in range(
default_multi_requests_causal_lm_batch.requests[0].stopping_criteria.max_new_tokens - default_multi_requests_causal_lm_batch.requests[1].stopping_criteria.max_new_tokens - 1 stopping_criterias[0].max_new_tokens - stopping_criterias[1].max_new_tokens - 1
): ):
generations, next_batch, _ = default_causal_lm.generate_token([next_batch]) generations, next_batch, _ = default_causal_lm.generate_token(next_batch)
assert len(generations) == len(next_batch) assert len(generations) == len(next_batch)
generations, next_batch, _ = default_causal_lm.generate_token([next_batch]) generations, next_batch, _ = default_causal_lm.generate_token(next_batch)
assert next_batch is None assert next_batch is None
assert len(generations) == 1 assert len(generations) == 1
assert generations[0].generated_text.text == "ing the effect of a new method for the detection" assert generations[0].generated_text.text == ".java:784) at net.minecraft."
assert ( assert (
generations[0].request_id generations[0].request_id
== default_multi_requests_causal_lm_batch.requests[0].data.id == default_multi_requests_causal_lm_batch.requests[0].id
) )
assert ( assert (
generations[0].generated_text.generated_tokens generations[0].generated_text.generated_tokens
== default_multi_requests_causal_lm_batch.requests[0].stopping_criteria.max_new_tokens == default_multi_requests_causal_lm_batch.stopping_criterias[0].max_new_tokens
) )
@ -232,13 +224,11 @@ def test_batch_concatenate(
default_causal_lm, default_causal_lm_batch, default_multi_requests_causal_lm_batch default_causal_lm, default_causal_lm_batch, default_multi_requests_causal_lm_batch
): ):
next_batch_0 = default_causal_lm_batch next_batch_0 = default_causal_lm_batch
_, next_batch_0, _ = default_causal_lm.generate_token([next_batch_0]) _, next_batch_0, _ = default_causal_lm.generate_token(next_batch_0)
_, next_batch_0, _ = default_causal_lm.generate_token([next_batch_0]) _, next_batch_0, _ = default_causal_lm.generate_token(next_batch_0)
_, next_batch_0, _ = default_causal_lm.generate_token([next_batch_0])
next_batch_1 = default_multi_requests_causal_lm_batch next_batch_1 = default_multi_requests_causal_lm_batch
_, next_batch_1, _ = default_causal_lm.generate_token([next_batch_1]) _, next_batch_1, _ = default_causal_lm.generate_token(next_batch_1)
_, next_batch_1, _ = default_causal_lm.generate_token([next_batch_1])
# Clone past_key_values before concatenating to compare after, # Clone past_key_values before concatenating to compare after,
# because they are removed from the concatenated batches # because they are removed from the concatenated batches
@ -251,135 +241,113 @@ def test_batch_concatenate(
next_batch = CausalLMBatch.concatenate([next_batch_0, next_batch_1]) next_batch = CausalLMBatch.concatenate([next_batch_0, next_batch_1])
assert torch.equal(next_batch.requests[0].all_input_ids, next_batch_0.requests[0].all_input_ids) assert torch.equal(next_batch.all_input_ids[0], next_batch_0.all_input_ids[0])
assert torch.equal(next_batch.requests[1].all_input_ids, next_batch_1.requests[0].all_input_ids) assert torch.equal(next_batch.all_input_ids[1], next_batch_1.all_input_ids[0])
assert torch.equal(next_batch.requests[2].all_input_ids, next_batch_1.requests[1].all_input_ids) assert torch.equal(next_batch.all_input_ids[2], next_batch_1.all_input_ids[1])
assert torch.all( assert torch.all(
next_batch.attention_mask[0:2, -next_batch.right_padding - 3: -next_batch.right_padding] == 1 next_batch.attention_mask[0, : -next_batch.padding_right_offset] == 1
) )
assert torch.all( assert torch.all(
next_batch.attention_mask[2, -next_batch.right_padding - 4: -next_batch.right_padding] == 1 next_batch.attention_mask[1:, 1 : -next_batch.padding_right_offset] == 1
) )
assert torch.all( assert torch.all(next_batch.attention_mask[1:, 3:] == 0)
next_batch.attention_mask[3, -next_batch.right_padding - 3: -next_batch.right_padding] == 1
)
assert torch.all(
next_batch.attention_mask[0:2, :-next_batch.right_padding-3] == 0)
assert torch.all(
next_batch.attention_mask[2, :-next_batch.right_padding-4] == 0)
assert torch.all(
next_batch.attention_mask[3, :-next_batch.right_padding-3] == 0)
assert next_batch.batch_id == 0 assert next_batch.batch_id == 0
assert next_batch.input_ids[0,-next_batch.right_padding - 3] == 1 assert next_batch.input_ids[0, 0] == 12355
assert next_batch.input_ids[0,-next_batch.right_padding - 2] == 4321 assert torch.all(next_batch.input_ids[1:] == 13)
assert next_batch.input_ids[0,-next_batch.right_padding - 1] == 292
assert next_batch.max_input_length == 129 assert next_batch.input_lengths == [3, 2, 2]
assert next_batch.max_input_length == 3
assert torch.all(next_batch.input_ids[0,-next_batch.right_padding:] == PAD_TOKEN)
assert torch.all(next_batch.input_ids[1,-next_batch.right_padding:] == PAD_TOKEN)
assert torch.all(next_batch.input_ids[2,-next_batch.right_padding:] == PAD_TOKEN)
assert torch.all(next_batch.input_ids[3,-next_batch.right_padding:] == PAD_TOKEN)
assert next_batch.input_length == PAD_SEQUENCE_TO_MULTIPLE_OF +1
assert next_batch.max_input_length == PAD_SEQUENCE_TO_MULTIPLE_OF +1
assert next_batch.requests[0] == next_batch_0.requests[0] assert next_batch.requests[0] == next_batch_0.requests[0]
assert next_batch.requests[1:] == next_batch_1.requests assert next_batch.requests[1:] == next_batch_1.requests
assert next_batch.requests[0].stopping_criteria == next_batch_0.requests[0].stopping_criteria assert next_batch.next_token_choosers[0] == next_batch_0.next_token_choosers[0]
assert next_batch.requests[1].stopping_criteria == next_batch_1.requests[0].stopping_criteria assert next_batch.next_token_choosers[1:] == next_batch_1.next_token_choosers
assert next_batch.requests[2].stopping_criteria == next_batch_1.requests[1].stopping_criteria
assert next_batch.stopping_criterias[0] == next_batch_0.stopping_criterias[0]
assert next_batch.past_key_values is not None assert next_batch.stopping_criterias[1:] == next_batch_1.stopping_criterias
assert all([p[0].shape == (8, 32, 2048, 128) for p in next_batch.past_key_values])
assert all([p[1].shape == (8, 32, 2048, 128) for p in next_batch.past_key_values])
assert next_batch.past_key_values is not None assert next_batch.past_key_values is not None
assert all([p[0].shape == (3, 12, 2, 64) for p in next_batch.past_key_values])
assert all([p[1].shape == (3, 12, 2, 64) for p in next_batch.past_key_values])
for i, past in enumerate(next_batch.past_key_values): for i, past in enumerate(next_batch.past_key_values):
assert torch.equal(next_batch_0_past_key_values[i][0][0, 0,0:128], past[0][0][0][1:129]) assert torch.equal(next_batch_0_past_key_values[i][0][0, :, -2:], past[0][0])
assert torch.equal( assert torch.equal(
next_batch_1_past_key_values[i][0][:, :, 0:1][0], past[0][1:, :, 1 :2, :][0] next_batch_1_past_key_values[i][0][:, :, -1:], past[0][1:, :, -1:, :]
) )
assert torch.equal(next_batch_0_past_key_values[i][1][0, 0,0:128], past[1][0][0][1:129]) assert torch.equal(next_batch_0_past_key_values[i][1][0, :, -2:], past[1][0])
assert torch.equal( assert torch.equal(
next_batch_1_past_key_values[i][1][:, :, 0:1][0], past[1][1:, :, 1 :2, :][0] next_batch_1_past_key_values[i][1][:, :, -1:], past[1][1:, :, -1:, :]
) )
generations, next_batch, _ = default_causal_lm.generate_token([next_batch])
for _ in range( for _ in range(
default_multi_requests_causal_lm_batch.requests[1].stopping_criteria.max_new_tokens - 2 default_multi_requests_causal_lm_batch.stopping_criterias[1].max_new_tokens - 2
): ):
generations, next_batch, _ = default_causal_lm.generate_token([next_batch]) generations, next_batch, _ = default_causal_lm.generate_token(next_batch)
assert len(generations) == len(next_batch) assert len(generations) == len(next_batch)
generations, next_batch, _ = default_causal_lm.generate_token([next_batch]) generations, next_batch, _ = default_causal_lm.generate_token(next_batch)
assert next_batch is not None assert next_batch is not None
assert len(generations) == 3 assert len(generations) == 3
assert generations[2].generated_text.text == "ing the effect of a" assert generations[2].generated_text.text == ".java:784)"
assert ( assert (
generations[2].request_id generations[2].request_id
== default_multi_requests_causal_lm_batch.requests[1].data.id == default_multi_requests_causal_lm_batch.requests[1].id
) )
assert ( assert (
generations[2].generated_text.generated_tokens generations[2].generated_text.generated_tokens
== default_multi_requests_causal_lm_batch.requests[1].stopping_criteria.max_new_tokens == default_multi_requests_causal_lm_batch.stopping_criterias[1].max_new_tokens
) )
next_batch = next_batch.filter( next_batch = next_batch.filter(
[next_batch.requests[0].data.id, next_batch.requests[1].data.id] [next_batch.requests[0].id, next_batch.requests[1].id]
) )
for _ in range( for _ in range(
default_causal_lm_batch.requests[0].stopping_criteria.max_new_tokens default_causal_lm_batch.stopping_criterias[0].max_new_tokens
- default_multi_requests_causal_lm_batch.requests[1].stopping_criteria.max_new_tokens - default_multi_requests_causal_lm_batch.stopping_criterias[1].max_new_tokens
- 2 - 2
): ):
generations, next_batch, _ = default_causal_lm.generate_token([next_batch]) generations, next_batch, _ = default_causal_lm.generate_token(next_batch)
assert len(generations) == len(next_batch) assert len(generations) == len(next_batch)
generations, next_batch, _ = default_causal_lm.generate_token([next_batch]) generations, next_batch, _ = default_causal_lm.generate_token(next_batch)
assert next_batch is not None assert next_batch is not None
assert len(generations) == 2 assert len(generations) == 2
assert generations[0].generated_text.text == "ing the effect of a new method for the detection" assert generations[0].generated_text.text == ".java:784) at net.minecraft."
assert generations[0].request_id == default_causal_lm_batch.requests[0].data.id assert generations[0].request_id == default_causal_lm_batch.requests[0].id
assert ( assert (
generations[0].generated_text.generated_tokens generations[0].generated_text.generated_tokens
== default_causal_lm_batch.requests[0].stopping_criteria.max_new_tokens == default_causal_lm_batch.stopping_criterias[0].max_new_tokens
) )
next_batch = next_batch.filter([next_batch.requests[1].data.id]) next_batch = next_batch.filter([next_batch.requests[1].id])
for _ in range( for _ in range(
default_multi_requests_causal_lm_batch.requests[0].stopping_criteria.max_new_tokens default_multi_requests_causal_lm_batch.stopping_criterias[0].max_new_tokens
- default_causal_lm_batch.requests[0].stopping_criteria.max_new_tokens - default_causal_lm_batch.stopping_criterias[0].max_new_tokens
- default_multi_requests_causal_lm_batch.requests[1].stopping_criteria.max_new_tokens - default_multi_requests_causal_lm_batch.stopping_criterias[1].max_new_tokens
- 4 - 4
): ):
generations, next_batch, _ = default_causal_lm.generate_token([next_batch]) generations, next_batch, _ = default_causal_lm.generate_token(next_batch)
assert len(generations) == len(next_batch) assert len(generations) == len(next_batch)
generations, next_batch, _ = default_causal_lm.generate_token([next_batch]) generations, next_batch, _ = default_causal_lm.generate_token(next_batch)
assert next_batch is None assert next_batch is None
assert len(generations) == 1 assert len(generations) == 1
assert generations[0].generated_text.text == "ing the effect of a new method for the detection" assert generations[0].generated_text.text == ".java:784) at net.minecraft."
assert ( assert (
generations[0].request_id generations[0].request_id
== default_multi_requests_causal_lm_batch.requests[0].data.id == default_multi_requests_causal_lm_batch.requests[0].id
) )
assert ( assert (
generations[0].generated_text.generated_tokens generations[0].generated_text.generated_tokens
== default_multi_requests_causal_lm_batch.requests[0].stopping_criteria.max_new_tokens == default_multi_requests_causal_lm_batch.stopping_criterias[0].max_new_tokens
) )

245
server/tests/models/test_grammar.py
View File

@ -1,245 +0,0 @@
# Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.
import json
import pytest
import torch
from copy import copy
from text_generation_server.pb import generate_pb2
from text_generation_server.models import get_model
from text_generation_server.models.causal_lm import (
CausalLMBatch,
PAD_SEQUENCE_TO_MULTIPLE_OF,
)
PAD_TOKEN=0
@pytest.fixture
def default_pb_grammar_parameters():
grammar_schema = {
"properties": {
"activity": {
"type": "string"
},
"animals": {
"items": {
"type":"string"
},
"type": "array"
}
},
"required": ["activity", "animals"]
}
return generate_pb2.NextTokenChooserParameters(
temperature=1.0,
repetition_penalty=1.3,
top_k=0,
top_p=1.0,
typical_p=1.0,
do_sample=False,
grammar_type=generate_pb2.GrammarType.GRAMMAR_TYPE_JSON,
grammar=json.dumps(grammar_schema).encode('utf-8'),
)
@pytest.fixture(scope="session")
def default_grammar_response():
return [
29912, 376, 29874, 312, 2068, 1115, 29871, 13, 29908, 29890,
638, 292, 613, 259, 376, 273, 3039, 29879, 1115,518, 1678,
376, 26169, 3284, 4117, 3284, 336, 617, 6150, 3108, 500, 2
]
@pytest.fixture(scope="session")
def default_causal_lm():
return get_model("meta-llama/Llama-2-7b-hf", None, None, None, None)
@pytest.fixture(scope="session")
def default_tokenizer(default_causal_lm):
default_causal_lm.tokenizer.pad_token_id = PAD_TOKEN
return default_causal_lm.tokenizer
@pytest.fixture
def default_pb_request(default_pb_parameters):
return generate_pb2.Request(
id=0,
inputs="Test",
prefill_logprobs=True,
truncate=PAD_SEQUENCE_TO_MULTIPLE_OF,
parameters=default_pb_parameters,
stopping_parameters=generate_pb2.StoppingCriteriaParameters(stop_sequences=[], max_new_tokens=10),
)
@pytest.fixture
def default_pb_grammar_request(default_pb_grammar_parameters):
return generate_pb2.Request(
id=1,
inputs=f"Please use the following JSON schema to generate the output: I saw a puppy a cat and a raccoon during my bike ride in the park",
prefill_logprobs=True,
truncate=PAD_SEQUENCE_TO_MULTIPLE_OF,
parameters=default_pb_grammar_parameters,
stopping_parameters=generate_pb2.StoppingCriteriaParameters(stop_sequences=[], max_new_tokens=50),
)
@pytest.fixture
def default_pb_batch(default_pb_request):
return generate_pb2.Batch(id=0, requests=[default_pb_request], size=1)
@pytest.fixture
def default_pb_grammar_batch(default_pb_grammar_request):
return generate_pb2.Batch(id=1, requests=[default_pb_grammar_request], size=1)
@pytest.fixture
def default_causal_lm_batch(default_pb_batch, default_tokenizer):
return CausalLMBatch.from_pb(
default_pb_batch, default_tokenizer, torch.float32, torch.device("hpu")
)
@pytest.fixture
def default_causal_lm_grammar_batch(default_pb_grammar_batch, default_tokenizer):
return CausalLMBatch.from_pb(
default_pb_grammar_batch, default_tokenizer, torch.float32, torch.device("hpu")
)
@pytest.fixture
def default_two_causal_lm_grammar_batches(default_pb_grammar_request, default_tokenizer):
req_0 = default_pb_grammar_request
req_0.id = 0
req_1 = copy(default_pb_grammar_request)
req_1.id = 1
batch_0 = generate_pb2.Batch(id=0, requests=[req_0], size=1)
batch_1 = generate_pb2.Batch(id=1, requests=[req_1], size=1)
return [
CausalLMBatch.from_pb(
b, default_tokenizer, torch.float32, torch.device("hpu")
) for b in [batch_0, batch_1]
]
def test_single_grammar_batch(
default_causal_lm, default_causal_lm_grammar_batch, default_grammar_response
):
counter = 0
batch = default_causal_lm_grammar_batch
# prefill request
generations, next_batch, _ = default_causal_lm.generate_token([batch])
# generate untill done
while next_batch is not None:
generations, next_batch, _ = default_causal_lm.generate_token([next_batch])
assert len(generations) == 1
assert generations[0].tokens.token_ids[0] == default_grammar_response[counter]
counter += 1
print(generations[0].generated_text.text)
def test_multi_grammar_batches(
default_causal_lm, default_two_causal_lm_grammar_batches, default_grammar_response
):
counter_0, counter_1 = 0, 0
batch_0, batch_1 = default_two_causal_lm_grammar_batches
# prefill first request
generations, next_batch, _ = default_causal_lm.generate_token([batch_0])
generations, next_batch, _ = default_causal_lm.generate_token([next_batch])
assert len(generations) == 1
assert generations[0].tokens.token_ids[0] == default_grammar_response[counter_0]
counter_0 += 1
# prefill second request
generations, next_batch_1, _ = default_causal_lm.generate_token([batch_1])
# concatenate and generate
generations, next_batch, _ = default_causal_lm.generate_token([next_batch, next_batch_1])
assert len(generations) == 2
assert generations[0].tokens.token_ids[0] == default_grammar_response[counter_0]
assert generations[1].tokens.token_ids[0] == default_grammar_response[counter_1]
counter_0 += 1
counter_1 += 1
# generate untill first request is done
while generations[0].generated_text is None:
generations, next_batch, _ = default_causal_lm.generate_token([next_batch])
assert len(generations) == 2
assert generations[0].tokens.token_ids[0] == default_grammar_response[counter_0]
assert generations[1].tokens.token_ids[0] == default_grammar_response[counter_1]
counter_0 += 1
counter_1 += 1
# filter finished request
response = generations[0].generated_text.text
next_batch = next_batch.filter([next_batch.requests[1].data.id])
# generate last tokens for second request
while next_batch is not None:
generations, next_batch, _ = default_causal_lm.generate_token([next_batch])
assert len(generations) == 1
assert generations[0].tokens.token_ids[0] == default_grammar_response[counter_1]
counter_1 += 1
assert response == generations[0].generated_text.text
def test_grammar_and_causal_batch(
default_causal_lm, default_causal_lm_grammar_batch, default_causal_lm_batch, default_grammar_response
):
counter = 0
generations, next_batch, _ = default_causal_lm.generate_token([default_causal_lm_grammar_batch])
generations, next_batch, _ = default_causal_lm.generate_token([next_batch])
assert len(generations) == 1
assert generations[0].tokens.token_ids[0] == default_grammar_response[counter]
counter += 1
generations, next_batch, _ = default_causal_lm.generate_token([next_batch])
assert len(generations) == 1
assert generations[0].tokens.token_ids[0] == default_grammar_response[counter]
counter += 1
# prefill second request
generations, next_batch_1, _ = default_causal_lm.generate_token([default_causal_lm_batch])
# concatenate and generate
generations, next_batch, _ = default_causal_lm.generate_token([next_batch, next_batch_1])
assert len(generations) == 2
assert generations[0].tokens.token_ids[0] == default_grammar_response[counter]
counter += 1
# generate untill second request is done
for _ in range(
next_batch.requests[1].stopping_criteria.max_new_tokens - 1
):
generations, next_batch, _ = default_causal_lm.generate_token([next_batch])
assert len(generations) == 2
assert generations[0].tokens.token_ids[0] == default_grammar_response[counter]
counter += 1
# filter finished request
assert len(generations) == 2
assert (
generations[1].request_id == next_batch.requests[1].data.id
)
assert (
generations[1].generated_text.generated_tokens == next_batch.requests[1].stopping_criteria.max_new_tokens
)
assert generations[1].generated_text.text == "ing the effect of a new method for the detection"
next_batch = next_batch.filter([next_batch.requests[0].data.id])
# generate untill done
while next_batch is not None:
generations, next_batch, _ = default_causal_lm.generate_token([next_batch])
assert len(generations) == 1
assert generations[0].tokens.token_ids[0] == default_grammar_response[counter]
counter += 1

7
server/tests/models/test_seq2seq_lm.py
View File

@ -61,7 +61,6 @@ def default_multi_requests_seq2seq_lm_batch(default_pb_request, mt0_small_tokeni
) )
@pytest.mark.skip("seq2seq model not enabled on HPU yet")
def test_batch_from_pb(default_pb_batch, default_seq2seq_lm_batch): def test_batch_from_pb(default_pb_batch, default_seq2seq_lm_batch):
batch = default_seq2seq_lm_batch batch = default_seq2seq_lm_batch
sequence_length = len(default_seq2seq_lm_batch.input_ids[0]) sequence_length = len(default_seq2seq_lm_batch.input_ids[0])
@ -93,18 +92,15 @@ def test_batch_from_pb(default_pb_batch, default_seq2seq_lm_batch):
assert batch.max_decoder_input_length == batch.decoder_input_lengths[0] assert batch.max_decoder_input_length == batch.decoder_input_lengths[0]
@pytest.mark.skip("seq2seq model not enabled on HPU yet")
def test_batch_concatenate_no_prefill(default_seq2seq_lm_batch): def test_batch_concatenate_no_prefill(default_seq2seq_lm_batch):
with pytest.raises(ValueError): with pytest.raises(ValueError):
Seq2SeqLMBatch.concatenate([default_seq2seq_lm_batch, default_seq2seq_lm_batch]) Seq2SeqLMBatch.concatenate([default_seq2seq_lm_batch, default_seq2seq_lm_batch])
@pytest.mark.skip("seq2seq model not enabled on HPU yet")
def test_seq2seq_lm_batch_type(default_seq2seq_lm): def test_seq2seq_lm_batch_type(default_seq2seq_lm):
assert default_seq2seq_lm.batch_type == Seq2SeqLMBatch assert default_seq2seq_lm.batch_type == Seq2SeqLMBatch
@pytest.mark.skip("seq2seq model not enabled on HPU yet")
def test_seq2seq_lm_generate_token(default_seq2seq_lm, default_seq2seq_lm_batch): def test_seq2seq_lm_generate_token(default_seq2seq_lm, default_seq2seq_lm_batch):
sequence_length = len(default_seq2seq_lm_batch.input_ids[0]) sequence_length = len(default_seq2seq_lm_batch.input_ids[0])
generations, next_batch, _ = default_seq2seq_lm.generate_token( generations, next_batch, _ = default_seq2seq_lm.generate_token(
@ -172,7 +168,6 @@ def test_seq2seq_lm_generate_token(default_seq2seq_lm, default_seq2seq_lm_batch)
assert generations[0].request_id == 0 assert generations[0].request_id == 0
@pytest.mark.skip("seq2seq model not enabled on HPU yet")
def test_seq2seq_lm_generate_token_completion( def test_seq2seq_lm_generate_token_completion(
default_seq2seq_lm, default_seq2seq_lm_batch default_seq2seq_lm, default_seq2seq_lm_batch
): ):
@ -190,7 +185,6 @@ def test_seq2seq_lm_generate_token_completion(
assert generations[0].generated_text.generated_tokens == 7 assert generations[0].generated_text.generated_tokens == 7
@pytest.mark.skip("seq2seq model not enabled on HPU yet")
def test_seq2seq_lm_generate_token_completion_multi( def test_seq2seq_lm_generate_token_completion_multi(
default_seq2seq_lm, default_multi_requests_seq2seq_lm_batch default_seq2seq_lm, default_multi_requests_seq2seq_lm_batch
): ):
@ -228,7 +222,6 @@ def test_seq2seq_lm_generate_token_completion_multi(
assert generations[0].generated_text.generated_tokens == 7 assert generations[0].generated_text.generated_tokens == 7
@pytest.mark.skip("seq2seq model not enabled on HPU yet")
def test_batch_concatenate( def test_batch_concatenate(
default_seq2seq_lm, default_seq2seq_lm,
default_seq2seq_lm_batch, default_seq2seq_lm_batch,

372
server/tests/models/test_starcoder.py
View File

@ -1,372 +0,0 @@
# Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.
import pytest
import torch
from copy import copy
from text_generation_server.pb import generate_pb2
from text_generation_server.models import get_model
from text_generation_server.models.starcoder import StarCoderCausalLMBatch
from text_generation_server.models.causal_lm import (
PREFILL_BATCH_BUCKET_SIZE,
PAD_SEQUENCE_TO_MULTIPLE_OF,
MAX_TOTAL_TOKENS,
BATCH_BUCKET_SIZE,
)
PAD_TOKEN=0
@pytest.fixture(scope="session")
def default_starcoder():
return get_model("bigcode/starcoder", None, None, None, None)
@pytest.fixture(scope="session")
def default_tokenizer(default_starcoder):
default_starcoder.tokenizer.pad_token_id = PAD_TOKEN
return default_starcoder.tokenizer
@pytest.fixture
def default_pb_request(default_pb_parameters, default_pb_stop_parameters):
return generate_pb2.Request(
id=0,
inputs="Test",
prefill_logprobs=True,
truncate=PAD_SEQUENCE_TO_MULTIPLE_OF,
parameters=default_pb_parameters,
stopping_parameters=default_pb_stop_parameters,
)
@pytest.fixture
def default_pb_batch(default_pb_request):
return generate_pb2.Batch(id=0, requests=[default_pb_request], size=1)
@pytest.fixture
def default_starcoder_batch(default_pb_batch, default_tokenizer):
return StarCoderCausalLMBatch.from_pb(
default_pb_batch, default_tokenizer, torch.float32, torch.device("hpu")
)
@pytest.fixture
def default_multi_requests_starcoder_batch(default_pb_request, default_tokenizer):
req_0 = copy(default_pb_request)
req_0.id = 1
req_1 = default_pb_request
req_1.id = 2
req_1.stopping_parameters.max_new_tokens = 5
batch_pb = generate_pb2.Batch(id=1, requests=[req_0, req_1], size=2)
return StarCoderCausalLMBatch.from_pb(
batch_pb, default_tokenizer, torch.float32, torch.device("hpu")
)
def test_starcoder_batch_type(default_starcoder):
assert default_starcoder.batch_type == StarCoderCausalLMBatch
def test_batch_from_pb(default_pb_batch, default_starcoder_batch):
batch = default_starcoder_batch
assert batch.batch_id == default_pb_batch.id
assert len(batch.requests) == len(default_pb_batch.requests)
for r in range(0,len(default_pb_batch.requests)):
assert batch.requests[r].data == default_pb_batch.requests[r]
# For Gaudi we are adding padding of multiplication of bucket size
size_of_padded_to_bucket = ((default_pb_batch.size + PREFILL_BATCH_BUCKET_SIZE - 1) // PREFILL_BATCH_BUCKET_SIZE) * PREFILL_BATCH_BUCKET_SIZE
assert len(batch.input_ids) == size_of_padded_to_bucket
assert batch.input_ids.shape == torch.Size([4, 128])
assert batch.input_ids[0][-2] == 1006
assert batch.input_ids[1][-2] == 49
assert batch.input_ids[2][-2] == 49
assert batch.attention_mask[0][-2] == 1
assert batch.attention_mask[1][-2] == 1
assert batch.attention_mask[2][-2] == 1
assert torch.all(batch.attention_mask[0, :-3] == 0)
assert batch.past_key_values is None
assert all(
[
torch.equal(input_ids, request.all_input_ids[:batch.input_length + 1, 0])
for input_ids, request in zip(batch.input_ids, batch.requests)
]
)
assert len(batch) == default_pb_batch.size
assert batch.max_input_length + 1 == default_pb_batch.requests[0].truncate
def test_starcoder_generate_token(default_starcoder, default_starcoder_batch):
sequence_length = len(default_starcoder_batch.requests[0].all_input_ids)
generations, next_batch, _ = default_starcoder.generate_token([default_starcoder_batch])
padding = next_batch.requests[0].stopping_criteria.max_new_tokens
assert isinstance(next_batch, StarCoderCausalLMBatch)
assert len(next_batch.attention_mask[0]) == PAD_SEQUENCE_TO_MULTIPLE_OF
assert next_batch.requests[0].all_input_ids[-padding-2] == 1006
assert torch.all(next_batch.requests[0].all_input_ids[-padding-1:] == PAD_TOKEN)
assert torch.all(next_batch.requests[0].all_input_ids[:-padding-3] == PAD_TOKEN)
generations, next_batch, _ = default_starcoder.generate_token([default_starcoder_batch])
assert torch.all(next_batch.attention_mask[0][PAD_SEQUENCE_TO_MULTIPLE_OF-2:PAD_SEQUENCE_TO_MULTIPLE_OF] == 1)
assert torch.all(next_batch.attention_mask[0][:PAD_SEQUENCE_TO_MULTIPLE_OF-3] == 0)
assert torch.all(next_batch.attention_mask[0][PAD_SEQUENCE_TO_MULTIPLE_OF+1:] == 0)
assert next_batch.requests[0].all_input_ids[-padding-2] == 1006
assert next_batch.requests[0].all_input_ids[-padding-1] == 26
assert torch.all(next_batch.requests[0].all_input_ids[-padding:] == PAD_TOKEN)
assert torch.all(next_batch.requests[0].all_input_ids[:-padding-3] == PAD_TOKEN)
assert next_batch.input_length == PAD_SEQUENCE_TO_MULTIPLE_OF
assert next_batch.max_input_length == next_batch.input_length
assert next_batch.past_key_values is not None
assert all(
[p[0].shape == (MAX_TOTAL_TOKENS, 256) for p in next_batch.past_key_values]
)
assert all(
[p[1].shape == (MAX_TOTAL_TOKENS, 256) for p in next_batch.past_key_values]
)
assert all([generation.generated_text is None for generation in generations])
assert all([len(generation.prefill_tokens) == PAD_SEQUENCE_TO_MULTIPLE_OF-1 for generation in generations])
assert all([generation.tokens.token_ids[0] == 26 for generation in generations])
assert all([generation.tokens.texts[0] == "(" for generation in generations])
assert generations[0].request_id == 0
def test_starcoder_generate_token_completion(
default_starcoder, default_starcoder_batch
):
next_batch = default_starcoder_batch
generations, next_batch, _ = default_starcoder.generate_token([next_batch])
for _ in range(default_starcoder_batch.requests[0].stopping_criteria.max_new_tokens - 1):
generations, next_batch, _ = default_starcoder.generate_token([next_batch])
assert len(generations) == len(next_batch)
generations, next_batch, _ = default_starcoder.generate_token([next_batch])
assert next_batch is None
assert len(generations) == 1
assert generations[0].generated_text.text == '(self):\n """\n Test that the test'
assert generations[0].request_id == default_starcoder_batch.requests[0].data.id
assert (
generations[0].generated_text.generated_tokens
== default_starcoder_batch.requests[0].stopping_criteria.max_new_tokens
)
def test_starcoder_generate_token_completion_multi(
default_starcoder, default_multi_requests_starcoder_batch
):
next_batch = default_multi_requests_starcoder_batch
generations, next_batch, _ = default_starcoder.generate_token([next_batch])
for i in range(
default_multi_requests_starcoder_batch.requests[1].stopping_criteria.max_new_tokens - 1
):
generations, next_batch, _ = default_starcoder.generate_token([next_batch])
assert len(generations) == len(next_batch)
generations, next_batch, _ = default_starcoder.generate_token([next_batch])
assert next_batch is not None
assert len(generations) == 2
assert generations[1].generated_text.text == '(self):\n """'
assert (
generations[1].request_id
== default_multi_requests_starcoder_batch.requests[1].data.id
)
assert (
generations[1].generated_text.generated_tokens
== default_multi_requests_starcoder_batch.requests[1].stopping_criteria.max_new_tokens
)
next_batch = next_batch.filter([next_batch.requests[0].data.id])
for _ in range(
default_multi_requests_starcoder_batch.requests[0].stopping_criteria.max_new_tokens - default_multi_requests_starcoder_batch.requests[1].stopping_criteria.max_new_tokens - 1
):
generations, next_batch, _ = default_starcoder.generate_token([next_batch])
assert len(generations) == len(next_batch)
generations, next_batch, _ = default_starcoder.generate_token([next_batch])
assert next_batch is None
assert len(generations) == 1
assert generations[0].generated_text.text == '(self):\n """\n Test that the test'
assert (
generations[0].request_id
== default_multi_requests_starcoder_batch.requests[0].data.id
)
assert (
generations[0].generated_text.generated_tokens
== default_multi_requests_starcoder_batch.requests[0].stopping_criteria.max_new_tokens
)
def test_batch_concatenate(
default_starcoder, default_starcoder_batch, default_multi_requests_starcoder_batch
):
next_batch_0 = default_starcoder_batch
_, next_batch_0, _ = default_starcoder.generate_token([next_batch_0])
_, next_batch_0, _ = default_starcoder.generate_token([next_batch_0])
_, next_batch_0, _ = default_starcoder.generate_token([next_batch_0])
next_batch_1 = default_multi_requests_starcoder_batch
_, next_batch_1, _ = default_starcoder.generate_token([next_batch_1])
_, next_batch_1, _ = default_starcoder.generate_token([next_batch_1])
# Clone past_key_values before concatenating to compare after,
# because they are removed from the concatenated batches
next_batch_0_past_key_values = [x.clone() for x in next_batch_0.past_key_values]
next_batch_1_past_key_values = [x.clone() for x in next_batch_1.past_key_values]
next_batch = StarCoderCausalLMBatch.concatenate([next_batch_0, next_batch_1])
assert torch.equal(next_batch.requests[0].all_input_ids, next_batch_0.requests[0].all_input_ids)
assert torch.equal(next_batch.requests[1].all_input_ids, next_batch_1.requests[0].all_input_ids)
assert torch.equal(next_batch.requests[2].all_input_ids, next_batch_1.requests[1].all_input_ids)
assert torch.all(
next_batch.attention_mask[0:2, -next_batch.right_padding - 2: -next_batch.right_padding] == 1
)
assert torch.all(
next_batch.attention_mask[2, -next_batch.right_padding - 3: -next_batch.right_padding] == 1
)
assert torch.all(
next_batch.attention_mask[3, -next_batch.right_padding - 2: -next_batch.right_padding] == 1
)
assert torch.all(
next_batch.attention_mask[0:2, :-next_batch.right_padding-2] == 0)
assert torch.all(
next_batch.attention_mask[2, :-next_batch.right_padding-4] == 0)
assert torch.all(
next_batch.attention_mask[3, :-next_batch.right_padding-3] == 0)
assert next_batch.batch_id == 0
assert next_batch.input_ids[0,-next_batch.right_padding - 2] == 1006
assert next_batch.input_ids[0,-next_batch.right_padding - 1] == 26
assert next_batch.max_input_length == 129
assert torch.all(next_batch.input_ids[0,-next_batch.right_padding:] == PAD_TOKEN)
assert torch.all(next_batch.input_ids[1,-next_batch.right_padding:] == PAD_TOKEN)
assert torch.all(next_batch.input_ids[2,-next_batch.right_padding:] == PAD_TOKEN)
assert torch.all(next_batch.input_ids[3,-next_batch.right_padding:] == PAD_TOKEN)
assert next_batch.input_length == PAD_SEQUENCE_TO_MULTIPLE_OF +1
assert next_batch.max_input_length == PAD_SEQUENCE_TO_MULTIPLE_OF + 1
assert next_batch.requests[0] == next_batch_0.requests[0]
assert next_batch.requests[1:] == next_batch_1.requests
assert next_batch.requests[0].stopping_criteria == next_batch_0.requests[0].stopping_criteria
assert next_batch.requests[1].stopping_criteria == next_batch_1.requests[0].stopping_criteria
assert next_batch.requests[2].stopping_criteria == next_batch_1.requests[1].stopping_criteria
assert next_batch.past_key_values is not None
assert all([p[0].shape == (2048, 256) for p in next_batch.past_key_values])
assert all([p[1].shape == (2048, 256) for p in next_batch.past_key_values])
assert next_batch.past_key_values is not None
for i, past in enumerate(next_batch.past_key_values):
assert torch.equal(next_batch_0_past_key_values[i][0,0,0:128], past[0][1:129][0, 0:128])
assert torch.equal(next_batch_0_past_key_values[i][0,1,0:128], past[1][1:129][0, 0:128])
assert torch.equal(
next_batch_1_past_key_values[i][:, :, 0:1][0][0][0], past[0][1:, :][0][0]
)
assert torch.equal(
next_batch_1_past_key_values[i][1:, :, 0:1][0][0][0], past[1][1:, :][0][0]
)
generations, next_batch, _ = default_starcoder.generate_token([next_batch])
for _ in range(
default_multi_requests_starcoder_batch.requests[1].stopping_criteria.max_new_tokens - 2
):
generations, next_batch, _ = default_starcoder.generate_token([next_batch])
assert len(generations) == len(next_batch)
generations, next_batch, _ = default_starcoder.generate_token([next_batch])
assert next_batch is not None
assert len(generations) == 3
assert generations[2].generated_text.text == '(self):\n """'
assert (
generations[2].request_id
== default_multi_requests_starcoder_batch.requests[1].data.id
)
assert (
generations[2].generated_text.generated_tokens
== default_multi_requests_starcoder_batch.requests[1].stopping_criteria.max_new_tokens
)
next_batch = next_batch.filter(
[next_batch.requests[0].data.id, next_batch.requests[1].data.id]
)
for _ in range(
default_starcoder_batch.requests[0].stopping_criteria.max_new_tokens
- default_multi_requests_starcoder_batch.requests[1].stopping_criteria.max_new_tokens
- 2
):
generations, next_batch, _ = default_starcoder.generate_token([next_batch])
assert len(generations) == len(next_batch)
generations, next_batch, _ = default_starcoder.generate_token([next_batch])
assert next_batch is not None
assert len(generations) == 2
assert generations[0].generated_text.text == '(self):\n """\n Test that the test'
assert generations[0].request_id == default_starcoder_batch.requests[0].data.id
assert (
generations[0].generated_text.generated_tokens
== default_starcoder_batch.requests[0].stopping_criteria.max_new_tokens
)
next_batch = next_batch.filter([next_batch.requests[1].data.id])
for _ in range(
default_multi_requests_starcoder_batch.requests[0].stopping_criteria.max_new_tokens
- default_starcoder_batch.requests[0].stopping_criteria.max_new_tokens
- default_multi_requests_starcoder_batch.requests[1].stopping_criteria.max_new_tokens
- 4
):
generations, next_batch, _ = default_starcoder.generate_token([next_batch])
assert len(generations) == len(next_batch)
generations, next_batch, _ = default_starcoder.generate_token([next_batch])
assert next_batch is None
assert len(generations) == 1
assert generations[0].generated_text.text == '(self):\n """\n Test that the test'
assert (
generations[0].request_id
== default_multi_requests_starcoder_batch.requests[0].data.id
)
assert (
generations[0].generated_text.generated_tokens
== default_multi_requests_starcoder_batch.requests[0].stopping_criteria.max_new_tokens
)

44
server/tests/utils/test_tokens.py
View File

@ -1,26 +1,12 @@
# Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.
import pytest
import torch import torch
from transformers import AutoTokenizer
from text_generation_server.utils.tokens import ( from text_generation_server.utils.tokens import (
StopSequenceCriteria, StopSequenceCriteria,
StoppingCriteria, StoppingCriteria,
FinishReason, FinishReason,
batch_top_tokens, batch_top_tokens,
make_tokenizer_optional,
) )
@pytest.fixture
def skip_tokenizer_env_var():
import os
os.environ["SKIP_TOKENIZER_IN_TGI"] = "true"
yield
del os.environ['SKIP_TOKENIZER_IN_TGI']
def test_stop_sequence_criteria(): def test_stop_sequence_criteria():
criteria = StopSequenceCriteria("/test;") criteria = StopSequenceCriteria("/test;")
@ -100,33 +86,3 @@ def test_batch_top_tokens():
assert topn_tok_logprobs[2] == [[-1, -2, -3, -3]] assert topn_tok_logprobs[2] == [[-1, -2, -3, -3]]
assert topn_tok_logprobs[3] == [[-1, -2, -3, -3]] assert topn_tok_logprobs[3] == [[-1, -2, -3, -3]]
assert topn_tok_logprobs[4] == [[-1, -2, -3, -3, -4]] assert topn_tok_logprobs[4] == [[-1, -2, -3, -3, -4]]
def test_pass_through_tokenizer(skip_tokenizer_env_var):
tokenizer = AutoTokenizer.from_pretrained(
'meta-llama/Llama-2-7b-chat-hf',
revision=None,
padding_side="left",
truncation_side="left",
)
tokenizer.pad_token_id = 2
make_tokenizer_optional(tokenizer)
input = ["1, 1724, 338, 6483, 6509, 29973", "?"]
tokenized_inputs = tokenizer(
input,
return_tensors="pt",
padding="max_length",
return_token_type_ids=False,
truncation=True,
max_length=1024,
)
assert tokenized_inputs['input_ids'].size() == torch.Size([2, 1024])
assert torch.equal(tokenized_inputs['input_ids'][0][1018:], torch.tensor([1, 1724, 338, 6483, 6509, 29973]))
assert torch.equal(tokenized_inputs['input_ids'][1][1023:], torch.tensor([tokenizer.pad_token_id]))
decoded_tokens = tokenizer.decode(tokenized_inputs["input_ids"][0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
assert decoded_tokens.split(',')[1018:] == ['1', '1724', '338', '6483', '6509', '29973']
if __name__ == "__main__":
test_pass_through_tokenizer()

85
server/tests/utils/test_watermark.py
View File

@ -1,8 +1,7 @@
# Copyright (C) 2024 Habana Labs, Ltd. an Intel Company. # test_watermark_logits_processor.py
import os import os
import numpy as np import numpy as np
import pytest
import torch import torch
from text_generation_server.utils.watermark import WatermarkLogitsProcessor from text_generation_server.utils.watermark import WatermarkLogitsProcessor
@ -11,71 +10,35 @@ GAMMA = os.getenv("WATERMARK_GAMMA", 0.5)
DELTA = os.getenv("WATERMARK_DELTA", 2.0) DELTA = os.getenv("WATERMARK_DELTA", 2.0)
@pytest.fixture def test_seed_rng():
def hpu_device(): input_ids = [101, 2036, 3731, 102, 2003, 103]
return torch.device("hpu") processor = WatermarkLogitsProcessor()
processor._seed_rng(input_ids)
@pytest.fixture
def input_ids_list():
return [101, 2036, 3731, 102, 2003, 103]
@pytest.fixture
def input_ids_tensor(hpu_device):
return torch.tensor(
[[101, 2036, 3731, 102, 2003, 103]],
dtype=torch.int64,
device=hpu_device
)
@pytest.fixture
def scores(hpu_device):
return torch.tensor(
[[0.5, 0.3, 0.2, 0.8], [0.1, 0.2, 0.7, 0.9]],
device=hpu_device
)
def test_seed_rng(input_ids_list, hpu_device):
processor = WatermarkLogitsProcessor(device=hpu_device)
processor._seed_rng(input_ids_list)
assert isinstance(processor.rng, torch.Generator) assert isinstance(processor.rng, torch.Generator)
def test_seed_rng_tensor(input_ids_tensor, hpu_device): def test_get_greenlist_ids():
processor = WatermarkLogitsProcessor(device=hpu_device) input_ids = [101, 2036, 3731, 102, 2003, 103]
processor._seed_rng(input_ids_tensor) processor = WatermarkLogitsProcessor()
assert isinstance(processor.rng, torch.Generator) result = processor._get_greenlist_ids(input_ids, 10, torch.device("cpu"))
def test_get_greenlist_ids(input_ids_list, hpu_device):
processor = WatermarkLogitsProcessor(device=hpu_device)
result = processor._get_greenlist_ids(input_ids_list, 10, hpu_device)
assert max(result) <= 10 assert max(result) <= 10
assert len(result) == int(10 * 0.5) assert len(result) == int(10 * 0.5)
def test_get_greenlist_ids_tensor(input_ids_tensor, hpu_device): def test_calc_greenlist_mask():
processor = WatermarkLogitsProcessor(device=hpu_device) processor = WatermarkLogitsProcessor()
result = processor._get_greenlist_ids(input_ids_tensor, 10, hpu_device) scores = torch.tensor([[0.5, 0.3, 0.2, 0.8], [0.1, 0.2, 0.7, 0.9]])
assert max(result) <= 10 greenlist_token_ids = torch.tensor([2, 3])
assert len(result) == int(10 * 0.5)
def test_calc_greenlist_mask(scores, hpu_device):
processor = WatermarkLogitsProcessor(device=hpu_device)
greenlist_token_ids = torch.tensor([2, 3], device=hpu_device)
result = processor._calc_greenlist_mask(scores, greenlist_token_ids) result = processor._calc_greenlist_mask(scores, greenlist_token_ids)
assert result.tolist() == [[False, False, False, False], [False, False, True, True]] assert result.tolist() == [[False, False, False, False], [False, False, True, True]]
assert result.shape == scores.shape assert result.shape == scores.shape
def test_bias_greenlist_logits(scores, hpu_device): def test_bias_greenlist_logits():
processor = WatermarkLogitsProcessor(device=hpu_device) processor = WatermarkLogitsProcessor()
scores = torch.tensor([[0.5, 0.3, 0.2, 0.8], [0.1, 0.2, 0.7, 0.9]])
green_tokens_mask = torch.tensor( green_tokens_mask = torch.tensor(
[[False, False, True, True], [False, False, False, True]], device=hpu_device [[False, False, True, True], [False, False, False, True]]
) )
greenlist_bias = 2.0 greenlist_bias = 2.0
result = processor._bias_greenlist_logits(scores, green_tokens_mask, greenlist_bias) result = processor._bias_greenlist_logits(scores, green_tokens_mask, greenlist_bias)
@ -83,13 +46,9 @@ def test_bias_greenlist_logits(scores, hpu_device):
assert result.shape == scores.shape assert result.shape == scores.shape
def test_call(input_ids_list, scores, hpu_device): def test_call():
processor = WatermarkLogitsProcessor(device=hpu_device) input_ids = [101, 2036, 3731, 102, 2003, 103]
result = processor(input_ids_list, scores) processor = WatermarkLogitsProcessor()
assert result.shape == scores.shape scores = torch.tensor([[0.5, 0.3, 0.2, 0.8], [0.1, 0.2, 0.7, 0.9]])
result = processor(input_ids, scores)
def test_call_tensor(input_ids_tensor, scores, hpu_device):
processor = WatermarkLogitsProcessor(device=hpu_device)
result = processor(input_ids_tensor, scores)
assert result.shape == scores.shape assert result.shape == scores.shape

105
server/text_generation_server/cli.py
View File

@ -1,8 +1,4 @@
# Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.
import os import os
import psutil
import signal
import sys import sys
import typer import typer
@ -18,6 +14,8 @@ app = typer.Typer()
class Quantization(str, Enum): class Quantization(str, Enum):
bitsandbytes = "bitsandbytes" bitsandbytes = "bitsandbytes"
bitsandbytes_nf4 = "bitsandbytes-nf4"
bitsandbytes_fp4 = "bitsandbytes-fp4"
gptq = "gptq" gptq = "gptq"
awq = "awq" awq = "awq"
eetq = "eetq" eetq = "eetq"
@ -44,6 +42,9 @@ def serve(
otlp_endpoint: Optional[str] = None, otlp_endpoint: Optional[str] = None,
): ):
if sharded: if sharded:
assert (
os.getenv("RANK", None) is not None
), "RANK must be set when sharded is True"
assert ( assert (
os.getenv("WORLD_SIZE", None) is not None os.getenv("WORLD_SIZE", None) is not None
), "WORLD_SIZE must be set when sharded is True" ), "WORLD_SIZE must be set when sharded is True"
@ -76,72 +77,25 @@ def serve(
# Downgrade enum into str for easier management later on # Downgrade enum into str for easier management later on
quantize = None if quantize is None else quantize.value quantize = None if quantize is None else quantize.value
dtype = "bfloat16" if dtype is None else dtype.value dtype = None if dtype is None else dtype.value
if dtype is not None and quantize not in {
logger.info("CLI SHARDED = {} DTYPE = {}".format(sharded, dtype)) None,
"bitsandbytes",
if sharded: "bitsandbytes-nf4",
tgi_file = Path(__file__).resolve().parent / "tgi_service.py" "bitsandbytes-fp4",
num_shard = int(os.getenv("WORLD_SIZE", "1")) }:
logger.info("CLI SHARDED = {}".format(num_shard)) raise RuntimeError(
import subprocess "Only 1 can be set between `dtype` and `quantize`, as they both decide how goes the final model."
)
cmd = f"deepspeed --num_nodes 1 --num_gpus {num_shard} --no_local_rank {tgi_file}"
cmd += f" --model_id {model_id} --revision {revision} --sharded {sharded}"
cmd += f" --dtype {dtype} --trust_remote_code {trust_remote_code} --uds_path {uds_path}"
if speculate is not None:
cmd += f"--speculate {speculate}"
logger.info("CLI server start deepspeed ={} ".format(cmd))
sys.stdout.flush()
sys.stderr.flush()
with subprocess.Popen(cmd, shell=True, executable="/bin/bash") as proc:
do_terminate = False
current_handler = signal.getsignal(signal.SIGTERM)
def terminate_handler(sig, frame):
nonlocal do_terminate
do_terminate = True
if callable(current_handler):
current_handler(sig, frame)
signal.signal(signal.SIGTERM, terminate_handler)
finished = False
while not finished:
try:
if do_terminate:
parent = psutil.Process(proc.pid)
all_procs = parent.children(recursive=True) + [parent]
for p in all_procs:
try:
p.terminate()
except psutil.NoSuchProcess:
pass
_, alive = psutil.wait_procs(all_procs, timeout=30)
for p in alive:
p.kill()
do_terminate = False
proc.wait(timeout=3)
except subprocess.TimeoutExpired:
pass
else:
finished = True
sys.stdout.flush()
sys.stderr.flush()
if proc.returncode != 0:
logger.error(f"{cmd} exited with status = {proc.returncode}")
return proc.returncode
else:
server.serve( server.serve(
model_id, model_id,
revision, revision,
sharded, sharded,
quantize,
speculate, speculate,
dtype, dtype,
trust_remote_code, trust_remote_code,
uds_path uds_path,
) )
@ -309,24 +263,30 @@ def download_weights(
) )
# Safetensors final filenames # Safetensors final filenames
local_st_files = [p.parent / f"{p.stem.lstrip('pytorch_')}.safetensors" for p in local_pt_files] local_st_files = [
p.parent / f"{p.stem.lstrip('pytorch_')}.safetensors"
for p in local_pt_files
]
try: try:
import transformers import transformers
from transformers import AutoConfig import json
config = AutoConfig.from_pretrained( if is_local_model:
model_id, config_filename = os.path.join(model_id, "config.json")
revision=revision, else:
config_filename = hf_hub_download(
model_id, revision=revision, filename="config.json"
) )
architecture = config.architectures[0] with open(config_filename, "r") as f:
config = json.load(f)
architecture = config["architectures"][0]
class_ = getattr(transformers, architecture) class_ = getattr(transformers, architecture)
# Name for this varible depends on transformers version. # Name for this varible depends on transformers version.
discard_names = getattr(class_, "_tied_weights_keys", []) discard_names = getattr(class_, "_tied_weights_keys", [])
discard_names.extend(getattr(class_, "_keys_to_ignore_on_load_missing", []))
except Exception: except Exception as e:
discard_names = [] discard_names = []
# Convert pytorch weights to safetensors # Convert pytorch weights to safetensors
utils.convert_files(local_pt_files, local_st_files, discard_names) utils.convert_files(local_pt_files, local_st_files, discard_names)
@ -344,6 +304,8 @@ def quantize(
percdamp: float = 0.01, percdamp: float = 0.01,
act_order: bool = False, act_order: bool = False,
): ):
if revision is None:
revision = "main"
download_weights( download_weights(
model_id=model_id, model_id=model_id,
revision=revision, revision=revision,
@ -357,6 +319,7 @@ def quantize(
bits=4, bits=4,
groupsize=128, groupsize=128,
output_dir=output_dir, output_dir=output_dir,
revision=revision,
trust_remote_code=trust_remote_code, trust_remote_code=trust_remote_code,
upload_to_model_id=upload_to_model_id, upload_to_model_id=upload_to_model_id,
percdamp=percdamp, percdamp=percdamp,

27
server/text_generation_server/habana_quantization_env.py
View File

@ -1,27 +0,0 @@
# Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.
import os
quant_config = os.getenv("QUANT_CONFIG", "")
is_quantization_enabled = quant_config != ""
if is_quantization_enabled:
os.environ.setdefault("ENABLE_EXPERIMENTAL_FLAGS", "true")
os.environ.setdefault("USE_DEFAULT_QUANT_PARAM", "true")
os.environ.setdefault("UPDATE_GRAPH_OUTPUT_MME", "false")
os.environ.setdefault("ENABLE_CALC_DYNAMIC_RANGE", "false")
os.environ.setdefault(
"UPDATE_MME_OUTPUT_PRECISION_FILTER", "v_proj,matmul_av")
os.environ.setdefault("EXPERIMENTAL_WEIGHT_SHARING", "FALSE")
def prepare_model_for_quantization(model):
if is_quantization_enabled:
if os.getenv("USE_INC", "1") != "0":
from neural_compressor.torch.quantization import FP8Config, convert
config = FP8Config.from_json_file(quant_config)
model = convert(model, config)
else:
import habana_quantization_toolkit
habana_quantization_toolkit.prep_model(model)
return model

9
server/text_generation_server/interceptor.py
View File

@ -1,5 +1,3 @@
# Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.
import torch import torch
import grpc import grpc
@ -8,8 +6,6 @@ from grpc_status import rpc_status
from grpc_interceptor.server import AsyncServerInterceptor from grpc_interceptor.server import AsyncServerInterceptor
from loguru import logger from loguru import logger
from typing import Callable, Any from typing import Callable, Any
import traceback
import os
class ExceptionInterceptor(AsyncServerInterceptor): class ExceptionInterceptor(AsyncServerInterceptor):
@ -24,7 +20,6 @@ class ExceptionInterceptor(AsyncServerInterceptor):
response = method(request_or_iterator, context) response = method(request_or_iterator, context)
return await response return await response
except Exception as err: except Exception as err:
trace = " " + traceback.format_exc() if os.environ.get('DUMP_STACK') else ''
method_name = method_name.split("/")[-1] method_name = method_name.split("/")[-1]
logger.exception(f"Method {method_name} encountered an error.") logger.exception(f"Method {method_name} encountered an error.")
@ -35,10 +30,8 @@ class ExceptionInterceptor(AsyncServerInterceptor):
if torch.cuda.is_available(): if torch.cuda.is_available():
torch.cuda.empty_cache() torch.cuda.empty_cache()
from .utils.debug import dbg_trace
dbg_trace('EXCEPTION', traceback.format_exc())
await context.abort_with_status( await context.abort_with_status(
rpc_status.to_status( rpc_status.to_status(
status_pb2.Status(code=code_pb2.INTERNAL, message=str(err) + trace) status_pb2.Status(code=code_pb2.INTERNAL, message=str(err))
) )
) )

2
server/text_generation_server/layers/gptq/exllamav2.py
View File

@ -119,6 +119,8 @@ def ext_make_q_matrix(w: dict, temp_dq, key: str = None):
none_tensor, none_tensor,
temp_dq, temp_dq,
) )
else:
RuntimeError("Cannot create handle")
DEVICE = None DEVICE = None

10
server/text_generation_server/layers/gptq/exllamav2.py.rej
View File

@ -1,10 +0,0 @@
diff a/server/text_generation_server/layers/gptq/exllamav2.py b/server/text_generation_server/layers/gptq/exllamav2.py (rejected hunks)
@@ -119,6 +119,8 @@ def ext_make_q_matrix(w: dict, temp_dq, key: str = None):
none_tensor,
temp_dq,
)
+ else:
+ RuntimeError("Cannot create handle")
DEVICE = None

783
server/text_generation_server/models/__init__.py
View File

@ -1,4 +1,5 @@
import torch import torch
import enum
import os import os
from loguru import logger from loguru import logger
@ -8,36 +9,269 @@ from huggingface_hub import hf_hub_download, HfApi
from typing import Optional from typing import Optional
from pathlib import Path from pathlib import Path
# Needed to properly setup habana_frameworks
import text_generation_server.habana_quantization_env as hq_env
from text_generation_server.utils.speculate import get_speculate, set_speculate from text_generation_server.utils.speculate import get_speculate, set_speculate
from text_generation_server.models.model import Model from text_generation_server.models.model import Model
from text_generation_server.models.causal_lm import CausalLM from text_generation_server.models.causal_lm import CausalLM
from text_generation_server.models.bloom import BLOOM from text_generation_server.models.flash_causal_lm import FlashCausalLM
from text_generation_server.models.starcoder import StarCoder from text_generation_server.models.bloom import BLOOMSharded
from text_generation_server.models.vlm_causal_lm import VlmCausalLM from text_generation_server.models.mpt import MPTSharded
from text_generation_server.models.custom_modeling.llava_next import ( from text_generation_server.models.seq2seq_lm import Seq2SeqLM
LlavaNextForConditionalGeneration, from text_generation_server.models.rw import RW
) from text_generation_server.models.opt import OPTSharded
from text_generation_server.models.galactica import GalacticaSharded
from text_generation_server.models.santacoder import SantaCoder
from text_generation_server.models.t5 import T5Sharded
from text_generation_server.models.gpt_neox import GPTNeoxSharded
from text_generation_server.models.phi import Phi
# The flag below controls whether to allow TF32 on matmul. This flag defaults to False
# in PyTorch 1.12 and later.
torch.backends.cuda.matmul.allow_tf32 = True
# The flag below controls whether to allow TF32 on cuDNN. This flag defaults to True.
from optimum.habana.transformers.modeling_utils import adapt_transformers_to_gaudi torch.backends.cudnn.allow_tf32 = True
# Disable gradients # Disable gradients
torch.set_grad_enabled(False) torch.set_grad_enabled(False)
__all__ = [
"Model",
"BLOOMSharded",
"CausalLM",
"GalacticaSharded",
"Seq2SeqLM",
"SantaCoder",
"OPTSharded",
"T5Sharded",
"get_model",
]
FLASH_ATT_ERROR_MESSAGE = "{} requires Flash Attention enabled models."
FLASH_ATTENTION = True
try:
from text_generation_server.models.flash_rw import FlashRWSharded
from text_generation_server.models.flash_gpt2 import FlashGPT2
from text_generation_server.models.flash_neox import FlashNeoXSharded
from text_generation_server.models.flash_llama import (
FlashLlama,
)
from text_generation_server.models.flash_qwen2 import (
FlashQwen2,
)
from text_generation_server.models.flash_cohere import (
FlashCohere,
)
from text_generation_server.models.flash_gemma import (
FlashGemma,
)
from text_generation_server.models.pali_gemma import (
PaliGemma,
)
from text_generation_server.models.flash_santacoder import (
FlashSantacoderSharded,
)
from text_generation_server.models.idefics import IDEFICSSharded
from text_generation_server.models.llava_next import LlavaNext
from text_generation_server.models.idefics2 import Idefics2
from text_generation_server.models.flash_mistral import FlashMistral
from text_generation_server.models.flash_mixtral import FlashMixtral
from text_generation_server.models.flash_phi import FlashPhi
from text_generation_server.models.flash_starcoder2 import FlashStarcoder2
from text_generation_server.models.flash_dbrx import FlashDbrx
from text_generation_server.utils.flash_attn import (
HAS_FLASH_ATTN_V2_CUDA,
HAS_FLASH_ATTN_V2_ROCM,
)
except ImportError as e:
logger.warning(f"Could not import Flash Attention enabled models: {e}")
FLASH_ATTENTION = False
HAS_FLASH_ATTN_V2_CUDA = False
HAS_FLASH_ATTN_V2_ROCM = False
if FLASH_ATTENTION:
__all__.append(FlashGPT2)
__all__.append(FlashNeoXSharded)
__all__.append(FlashRWSharded)
__all__.append(FlashSantacoderSharded)
__all__.append(FlashLlama)
__all__.append(IDEFICSSharded)
__all__.append(FlashMistral)
__all__.append(FlashMixtral)
__all__.append(FlashDbrx)
__all__.append(FlashPhi)
__all__.append(FlashQwen2)
__all__.append(FlashStarcoder2)
__all__.append(FlashGemma)
__all__.append(FlashCohere)
MAMBA_AVAILABLE = True
try:
from text_generation_server.models.mamba import Mamba
except ImportError as e:
logger.warning(f"Could not import Mamba: {e}")
MAMBA_AVAILABLE = False
if MAMBA_AVAILABLE:
__all__.append(Mamba)
class ModelType(enum.Enum):
IDEFICS2 = {
"type": "idefics2",
"name": "Idefics 2",
"url": "https://huggingface.co/HuggingFaceM4/idefics2-8b",
"multimodal": True,
}
LLAVA_NEXT = {
"type": "llava_next",
"name": "Llava Next (1.6)",
"url": "https://huggingface.co/llava-hf/llava-v1.6-vicuna-13b-hf",
"multimodal": True,
}
LLAMA = {
"type": "llama",
"name": "Llama",
"url": "https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct",
}
PHI3 = {
"type": "phi3",
"name": "Phi 3",
"url": "https://huggingface.co/microsoft/Phi-3-mini-4k-instruct",
}
GEMMA = {
"type": "gemma",
"name": "Gemma",
"url": "https://huggingface.co/google/gemma-7b",
}
COHERE = {
"type": "cohere",
"name": "Cohere",
"url": "https://huggingface.co/CohereForAI/c4ai-command-r-plus",
}
DBRX = {
"type": "dbrx",
"name": "Dbrx",
"url": "https://huggingface.co/databricks/dbrx-instruct",
}
MAMBA = {
"type": "ssm",
"name": "Mamba",
"url": "https://huggingface.co/state-spaces/mamba-2.8b-slimpj",
}
MISTRAL = {
"type": "mistral",
"name": "Mistral",
"url": "https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2",
}
MIXTRAL = {
"type": "mixtral",
"name": "Mixtral",
"url": "https://huggingface.co/mistralai/Mixtral-8x22B-Instruct-v0.1",
}
GPT_BIGCODE = {
"type": "gpt_bigcode",
"name": "Gpt Bigcode",
"url": "https://huggingface.co/bigcode/gpt_bigcode-santacoder",
}
PHI = {
"type": "phi",
"name": "Phi",
"url": "https://huggingface.co/microsoft/phi-1_5",
}
BAICHUAN = {
"type": "baichuan",
"name": "Baichuan",
"url": "https://huggingface.co/baichuan-inc/Baichuan2-7B-Chat",
}
FALCON = {
"type": "falcon",
"name": "Falcon",
"url": "https://huggingface.co/tiiuae/falcon-7b-instruct",
}
STARCODER2 = {
"type": "starcoder2",
"name": "StarCoder 2",
"url": "https://huggingface.co/bigcode/starcoder2-15b-instruct-v0.1",
}
QWEN2 = {
"type": "qwen2",
"name": "Qwen 2",
"url": "https://huggingface.co/bigcode/starcoder2-15b-instruct-v0.1",
}
OPT = {
"type": "opt",
"name": "Opt",
"url": "https://huggingface.co/facebook/opt-6.7b",
}
T5 = {
"type": "t5",
"name": "T5",
"url": "https://huggingface.co/google/flan-t5-xxl",
}
GALACTICA = {
"type": "galactica",
"name": "Galactica",
"url": "https://huggingface.co/facebook/galactica-120b",
}
SANTACODER = {
"type": "santacoder",
"name": "SantaCoder",
"url": "https://huggingface.co/bigcode/santacoder",
}
BLOOM = {
"type": "bloom",
"name": "Bloom",
"url": "https://huggingface.co/bigscience/bloom-560m",
}
MPT = {
"type": "mpt",
"name": "Mpt",
"url": "https://huggingface.co/mosaicml/mpt-7b-instruct",
}
GPT2 = {
"type": "gpt2",
"name": "Gpt2",
"url": "https://huggingface.co/openai-community/gpt2",
}
GPT_NEOX = {
"type": "gpt_neox",
"name": "Gpt Neox",
"url": "https://huggingface.co/EleutherAI/gpt-neox-20b",
}
IDEFICS = {
"type": "idefics",
"name": "Idefics",
"url": "https://huggingface.co/HuggingFaceM4/idefics-9b",
"multimodal": True,
}
__GLOBALS = locals()
for data in ModelType:
__GLOBALS[data.name] = data.value["type"]
def get_model( def get_model(
model_id: str, model_id: str,
revision: Optional[str], revision: Optional[str],
sharded: bool,
quantize: Optional[str],
speculate: Optional[int], speculate: Optional[int],
dtype: Optional[torch.dtype], dtype: Optional[str],
trust_remote_code: bool, trust_remote_code: bool,
) -> Model: ) -> Model:
adapt_transformers_to_gaudi() if dtype is None:
# Keep it as default for now and let
# every model resolve their own default dtype.
dtype = None
elif dtype == "float16":
dtype = torch.float16
elif dtype == "bfloat16":
dtype = torch.bfloat16
else:
raise RuntimeError(f"Unknown dtype {dtype}")
if speculate is not None: if speculate is not None:
set_speculate(speculate) set_speculate(speculate)
@ -152,35 +386,532 @@ def get_model(
if speculate > 0: if speculate > 0:
logger.info(f"Using speculation {method} with {speculate} input ids.") logger.info(f"Using speculation {method} with {speculate} input ids.")
model_type = config_dict["model_type"] if model_type is None:
# TODO: fix how we determine model type for Mamba
if "ssm_cfg" in config_dict:
# *only happens in Mamba case
model_type = "ssm"
else:
raise RuntimeError(
f"Could not determine model type for {model_id} revision {revision}"
)
quantization_config = config_dict.get("quantization_config", None)
if quantization_config is not None and quantize is None:
method = quantization_config.get("quant_method", None)
if method in {"gptq", "awq"}:
logger.info(f"Auto selecting quantization method {method}")
quantize = method
else:
logger.info(f"Unknown quantization method {method}")
if model_type == "gpt_bigcode": if model_type == MAMBA:
return StarCoder(model_id, revision, dtype) return Mamba(
if model_type == "bloom":
return BLOOM(
model_id, model_id,
revision, revision,
quantize=quantize,
speculator=speculator, speculator=speculator,
dtype=dtype, dtype=dtype,
trust_remote_code=trust_remote_code, trust_remote_code=trust_remote_code,
) )
if model_type == "llava_next": if model_id.startswith("facebook/galactica"):
return VlmCausalLM( return GalacticaSharded(
model_class=LlavaNextForConditionalGeneration, model_id,
model_id=model_id, revision,
revision=revision, quantize=quantize,
quantize=None,
speculator=speculator, speculator=speculator,
dtype=dtype, dtype=dtype,
trust_remote_code=trust_remote_code, trust_remote_code=trust_remote_code,
) )
if (
model_type == GPT_BIGCODE
or model_type == GPT2
and model_id.startswith("bigcode/")
):
if FLASH_ATTENTION:
return FlashSantacoderSharded(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
elif sharded:
raise NotImplementedError(
FLASH_ATT_ERROR_MESSAGE.format("Sharded Santacoder")
)
else:
return SantaCoder(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
if model_type == BLOOM:
return BLOOMSharded(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
elif model_type == MPT:
return MPTSharded(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
elif model_type == GPT2:
if FLASH_ATTENTION:
try:
return FlashGPT2(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
except RuntimeError as e:
# Lots of legacy models with various weight names.
logger.warning(f"Couldn't load flash gpt2 variant: {e}")
return CausalLM(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
elif sharded:
raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded GPT-2"))
else:
return CausalLM(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
elif model_type == GPT_NEOX:
if FLASH_ATTENTION:
return FlashNeoXSharded(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
elif sharded:
return GPTNeoxSharded(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
else:
return CausalLM(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
elif model_type == PHI:
if FLASH_ATTENTION:
return FlashPhi(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
else:
return CausalLM(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
elif model_type == "phi-msft":
if FLASH_ATTENTION:
raise NotImplementedError(
"Legacy phi-msft is not supported with Flash Attention"
)
else:
return Phi(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
elif model_type == LLAMA or model_type == BAICHUAN or model_type == PHI3:
if FLASH_ATTENTION:
return FlashLlama(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
elif sharded:
raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded Llama"))
else:
return CausalLM(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
if model_type == GEMMA:
if FLASH_ATTENTION:
return FlashGemma(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
elif sharded:
raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded Gemma"))
else:
return CausalLM(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
if model_type == COHERE:
if FLASH_ATTENTION:
return FlashCohere(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
elif sharded:
raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded Cohere"))
else:
return CausalLM(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
if model_type == DBRX:
if FLASH_ATTENTION:
return FlashDbrx(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
elif sharded:
raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded DBRX"))
else:
return CausalLM(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
if model_type in ["RefinedWeb", "RefinedWebModel", FALCON]:
if sharded:
if FLASH_ATTENTION:
if config_dict.get("alibi", False):
raise NotImplementedError("sharded is not supported for this model")
return FlashRWSharded(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format(f"Sharded Falcon"))
else:
if FLASH_ATTENTION and not config_dict.get("alibi", False):
return FlashRWSharded(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
else:
return RW(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
if model_type == MISTRAL:
sliding_window = config_dict.get("sliding_window", -1)
if (
((sliding_window is None or sliding_window == -1) and FLASH_ATTENTION)
or HAS_FLASH_ATTN_V2_CUDA
or HAS_FLASH_ATTN_V2_ROCM
):
return FlashMistral(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
elif sharded:
raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded Mistral"))
else:
return CausalLM(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
if model_type == MIXTRAL:
sliding_window = config_dict.get("sliding_window", -1)
if (
((sliding_window is None or sliding_window == -1) and FLASH_ATTENTION)
or HAS_FLASH_ATTN_V2_CUDA
or HAS_FLASH_ATTN_V2_ROCM
):
return FlashMixtral(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
elif sharded:
raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded Mixtral"))
else:
return CausalLM(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
if model_type == STARCODER2:
sliding_window = config_dict.get("sliding_window", -1)
if (
((sliding_window is None or sliding_window == -1) and FLASH_ATTENTION)
or HAS_FLASH_ATTN_V2_CUDA
or HAS_FLASH_ATTN_V2_ROCM
):
return FlashStarcoder2(
model_id,
revision,
quantize=quantize,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
elif sharded:
raise NotImplementedError(
FLASH_ATT_ERROR_MESSAGE.format("Sharded Starcoder2")
)
else:
return CausalLM(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
if model_type == QWEN2:
sliding_window = config_dict.get("sliding_window", -1)
if (
((sliding_window is None or sliding_window == -1) and FLASH_ATTENTION)
or HAS_FLASH_ATTN_V2_CUDA
or HAS_FLASH_ATTN_V2_ROCM
):
return FlashQwen2(
model_id,
revision,
quantize=quantize,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
elif sharded:
raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded Qwen2"))
else:
return CausalLM(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
if model_type == OPT:
return OPTSharded(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
if model_type == T5:
return T5Sharded(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
if model_type == IDEFICS:
if FLASH_ATTENTION:
return IDEFICSSharded(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
else:
raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Idefics"))
if model_type == IDEFICS2:
if FLASH_ATTENTION:
return Idefics2(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
else:
raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Idefics"))
if model_type == "paligemma":
if FLASH_ATTENTION:
return PaliGemma(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
else:
raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Idefics"))
if model_type == LLAVA_NEXT:
if FLASH_ATTENTION:
return LlavaNext(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
else:
raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("LlavaNext"))
if sharded:
raise NotImplementedError("sharded is not supported for AutoModel")
if quantize == "gptq":
raise NotImplementedError(
"gptq quantization is not supported for AutoModel, you can try to quantize it with `text-generation-server quantize ORIGINAL_MODEL_ID NEW_MODEL_ID`"
)
if quantize == "awq":
raise NotImplementedError("awq quantization is not supported for AutoModel")
elif (quantize == "bitsandbytes-fp4") or (quantize == "bitsandbytes-nf4"):
raise NotImplementedError("4bit quantization is not supported for AutoModel")
elif quantize == "eetq":
raise NotImplementedError("Eetq quantization is not supported for AutoModel")
if model_type in modeling_auto.MODEL_FOR_CAUSAL_LM_MAPPING_NAMES: if model_type in modeling_auto.MODEL_FOR_CAUSAL_LM_MAPPING_NAMES:
return CausalLM( return CausalLM(
model_id, model_id,
revision, revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
if model_type in modeling_auto.MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES:
return Seq2SeqLM(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
auto_map = config_dict.get("auto_map", None)
if trust_remote_code and auto_map is not None:
if "AutoModelForCausalLM" in auto_map.keys():
return CausalLM(
model_id,
revision,
quantize=quantize,
speculator=speculator,
dtype=dtype,
trust_remote_code=trust_remote_code,
)
if "AutoModelForSeq2SeqLM" in auto_map.keys():
return Seq2SeqLM(
model_id,
revision,
quantize=quantize,
speculator=speculator, speculator=speculator,
dtype=dtype, dtype=dtype,
trust_remote_code=trust_remote_code, trust_remote_code=trust_remote_code,

94
server/text_generation_server/models/bloom.py
View File

@ -1,14 +1,25 @@
# Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.
import torch import torch
import torch.distributed
from typing import Optional, Type from typing import Optional, Type
from transformers import PreTrainedTokenizerBase from transformers import (
AutoTokenizer,
AutoConfig,
PreTrainedTokenizerBase,
)
from text_generation_server.models.custom_modeling.bloom_modeling import (
BloomForCausalLM,
)
from text_generation_server.models import CausalLM from text_generation_server.models import CausalLM
from text_generation_server.models.causal_lm import CausalLMBatch from text_generation_server.models.causal_lm import CausalLMBatch
from text_generation_server.pb import generate_pb2 from text_generation_server.pb import generate_pb2
from text_generation_server.utils import (
initialize_torch_distributed,
weight_files,
Weights,
)
class BloomCausalLMBatch(CausalLMBatch): class BloomCausalLMBatch(CausalLMBatch):
@ -20,33 +31,88 @@ class BloomCausalLMBatch(CausalLMBatch):
dtype: torch.dtype, dtype: torch.dtype,
device: torch.device, device: torch.device,
) -> "CausalLMBatch": ) -> "CausalLMBatch":
batch = super().from_pb( batch = super().from_pb(pb=pb, tokenizer=tokenizer, dtype=dtype, device=device)
pb=pb,
tokenizer=tokenizer,
dtype=dtype,
device=device,
)
batch.keys_head_dim_last = False batch.keys_head_dim_last = False
return batch return batch
class BLOOM(CausalLM): class BLOOMSharded(CausalLM):
def __init__( def __init__(
self, self,
model_id: str, model_id: str,
revision: Optional[str] = None, revision: Optional[str] = None,
quantize: Optional[str] = None,
speculator: Optional[str] = None, speculator: Optional[str] = None,
dtype: Optional[torch.dtype] = None, dtype: Optional[torch.dtype] = None,
trust_remote_code: bool = False, trust_remote_code: bool = False,
): ):
super(BLOOM, self).__init__(
model_id=model_id, self.process_group, rank, world_size = initialize_torch_distributed()
if torch.cuda.is_available():
device = torch.device(f"cuda:{rank}")
dtype = torch.float16 if dtype is None else dtype
else:
device = torch.device("cpu")
dtype = torch.float32 if dtype is None else dtype
tokenizer = AutoTokenizer.from_pretrained(
model_id,
revision=revision, revision=revision,
speculator=speculator, padding_side="left",
dtype=dtype, truncation_side="left",
trust_remote_code=trust_remote_code, trust_remote_code=trust_remote_code,
) )
config = AutoConfig.from_pretrained(
model_id,
revision=revision,
slow_but_exact=False,
tp_parallel=True,
trust_remote_code=trust_remote_code,
)
config.pad_token_id = 3
config.quantize = quantize
config.speculator = speculator
torch.distributed.barrier(group=self.process_group)
filenames = weight_files(model_id, revision=revision, extension=".safetensors")
weights = Weights(
filenames,
device=device,
dtype=dtype,
process_group=self.process_group,
prefix="transformer",
)
if config.quantize == "gptq":
weights._set_gptq_params(model_id, revision)
model = BloomForCausalLM(config, weights)
torch.distributed.barrier(group=self.process_group)
super(CausalLM, self).__init__(
model=model,
tokenizer=tokenizer,
requires_padding=True,
dtype=dtype,
device=device,
rank=rank,
world_size=world_size,
)
@property @property
def batch_type(self) -> Type[CausalLMBatch]: def batch_type(self) -> Type[CausalLMBatch]:
return BloomCausalLMBatch return BloomCausalLMBatch
def forward(
self, input_ids, attention_mask, position_ids, past_key_values: Optional = None
):
outputs, speculative_logits = self.model.forward(
input_ids=input_ids,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
use_cache=True,
)
logits = outputs.logits
return logits, speculative_logits, outputs.past_key_values

1474
server/text_generation_server/models/causal_lm.py
View File

File diff suppressed because it is too large Load Diff

349
server/text_generation_server/models/custom_modeling/llava_next.py
View File

@ -21,12 +21,18 @@ import torch.utils.checkpoint
from torch import nn from torch import nn
from transformers.activations import ACT2FN from transformers.activations import ACT2FN
from transformers.models.llava_next.modeling_llava_next import (
unpad_image,
)
from optimum.habana.transformers.models import GaudiLlavaNextForConditionalGeneration
from transformers.image_processing_utils import select_best_resolution from transformers.image_processing_utils import select_best_resolution
from text_generation_server.models.custom_modeling.vlm import (
load_text_model,
load_vision_model,
)
from text_generation_server.layers import (
TensorParallelColumnLinear,
TensorParallelRowLinear,
)
def get_anyres_image_grid_shape(image_size, grid_pinpoints, patch_size): def get_anyres_image_grid_shape(image_size, grid_pinpoints, patch_size):
""" """
Calculate the shape of the image patch grid after the preprocessing for images of any resolution. Calculate the shape of the image patch grid after the preprocessing for images of any resolution.
@ -50,13 +56,100 @@ def get_anyres_image_grid_shape(image_size, grid_pinpoints, patch_size):
return height // patch_size, width // patch_size return height // patch_size, width // patch_size
class LlavaNextForConditionalGeneration(GaudiLlavaNextForConditionalGeneration): def unpad_image(tensor, original_size):
"""
Unpads a PyTorch tensor of a padded and resized image.
Args:
tensor (`torch.Tensor`):
The image tensor, assumed to be of shape (num_channels, height, width).
original_size (`tuple`):
The original size of the image (height, width).
Returns:
`torch.Tensor`: The unpadded image tensor.
"""
original_height, original_width = original_size
current_height, current_width = tensor.shape[1:]
original_aspect_ratio = original_width / original_height
current_aspect_ratio = current_width / current_height
if original_aspect_ratio > current_aspect_ratio:
scale_factor = current_width / original_width
new_height = int(original_height * scale_factor)
padding = (current_height - new_height) // 2
unpadded_tensor = tensor[:, padding : current_height - padding, :]
else:
scale_factor = current_height / original_height
new_width = int(original_width * scale_factor)
padding = (current_width - new_width) // 2
unpadded_tensor = tensor[:, :, padding : current_width - padding]
return unpadded_tensor
# Copied from transformers.models.llava.modeling_llava.LlavaMultiModalProjector with Llava->LlavaNext
class LlavaNextMultiModalProjector(nn.Module):
def __init__(self, prefix, config, weights):
super().__init__()
self.linear_1 = TensorParallelColumnLinear.load(
prefix=f"{prefix}.linear_1", config=config, weights=weights, bias=True
)
self.act = ACT2FN[config.projector_hidden_act]
self.linear_2 = TensorParallelRowLinear.load(
prefix=f"{prefix}.linear_2", config=config, weights=weights, bias=True
)
def forward(self, image_features):
hidden_states = self.linear_1(image_features)
hidden_states = self.act(hidden_states)
hidden_states = self.linear_2(hidden_states)
return hidden_states
class LlavaNextForConditionalGeneration(nn.Module):
def __init__(self, prefix, config, weights):
super().__init__()
config.vision_config.quantize = config.quantize
vision_config = config.vision_config
# Instead of selecting in hidden_states[-2].
# Instead compute only the n -2 + 1 layers and don't pool
if config.vision_feature_layer < 0:
vision_config.num_hidden_layers += config.vision_feature_layer + 1
else:
vision_config.num_hidden_layers = config.vision_feature_layer + 1
self.vision_tower = load_vision_model(
prefix="vision_tower" if not prefix else f"{prefix}.vision_tower",
config=config.vision_config,
weights=weights,
)
self.multi_modal_projector = LlavaNextMultiModalProjector(
prefix="multi_modal_projector", config=config, weights=weights
)
self.image_newline = weights.get_tensor("image_newline")
self.vocab_size = config.text_config.vocab_size
self.config = config
config.text_config.quantize = config.quantize
config.text_config.speculator = config.speculator
self.language_model = load_text_model(
prefix="language_model" if not prefix else f"{prefix}.language_model",
config=config.text_config,
weights=weights,
)
self.pad_token_id = (
config.pad_token_id if config.pad_token_id is not None else -1
)
def _merge_input_ids_with_image_features( def _merge_input_ids_with_image_features(
self, self,
input_ids: torch.Tensor,
inputs_embeds: torch.Tensor, inputs_embeds: torch.Tensor,
image_features: torch.Tensor, image_features: torch.Tensor,
input_ids: torch.Tensor,
): ):
"""In place merges in vision_embeddings with inputs_embeds.""" """In place merges in vision_embeddings with inputs_embeds."""
mask = input_ids == self.config.image_token_index mask = input_ids == self.config.image_token_index
@ -71,130 +164,60 @@ class LlavaNextForConditionalGeneration(GaudiLlavaNextForConditionalGeneration):
def forward( def forward(
self, self,
input_ids: torch.LongTensor = None, input_ids: torch.Tensor,
position_ids: torch.Tensor,
cu_seqlen_prefill: Optional[torch.Tensor],
kv_cache: List[Tuple[torch.Tensor, torch.Tensor]],
block_tables: torch.Tensor,
slots: torch.Tensor,
input_lengths: torch.Tensor,
max_s: int,
prefill_cache_indices: Optional[torch.Tensor],
lm_head_indices: Optional[torch.Tensor] = None,
pixel_values: torch.FloatTensor = None, pixel_values: torch.FloatTensor = None,
# Unused for this model
pixel_attention_mask=None,
image_sizes: Optional[torch.LongTensor] = None, image_sizes: Optional[torch.LongTensor] = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_values: Optional[List[torch.FloatTensor]] = None,
inputs_embeds: Optional[torch.FloatTensor] = None,
vision_feature_layer: Optional[int] = None,
vision_feature_select_strategy: Optional[str] = None,
labels: Optional[torch.LongTensor] = None,
use_cache: Optional[bool] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
token_idx: Optional[torch.Tensor] = None,
use_flash_attention: Optional[bool] = False,
flash_attention_recompute: Optional[bool] = False,
): ):
inputs_embeds = self.language_model.embed_tokens(input_ids)
if token_idx is not None: if pixel_values is not None and len(pixel_values) > 0:
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions # num_special_image_tokens = (input_ids == self.config.image_token_index).sum()
output_hidden_states = ( # assert num_special_image_tokens == len(pixel_values), f"Received {num_special_image_tokens} for {len(pixel_values)} images, this is invalid"
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
)
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
if inputs_embeds is None:
inputs_embeds = self.get_input_embeddings()(input_ids)
outputs = self.language_model(
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
inputs_embeds=inputs_embeds,
use_cache=use_cache,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
token_idx=token_idx,
use_flash_attention=use_flash_attention,
flash_attention_recompute=flash_attention_recompute,
)
logits = outputs[0]
if not return_dict:
output = (logits,) + outputs[1:]
return output
return outputs
def prepare_inputs_for_generation(
self,
input_ids,
past_key_values=None,
inputs_embeds=None,
pixel_values=None,
image_sizes=None,
attention_mask=None,
**kwargs,
):
"""
Inherits from LlavaForConditionalGeneration: https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/models/llava_next/modeling_llava_next.py#L635
The only differences are:
- add new args token_idx
- add the process of merging images into inputs_embeds
"""
token_idx = kwargs.get("token_idx", None)
if token_idx is None:
return super().prepare_inputs_for_generation(
input_ids=input_ids,
past_key_values=past_key_values,
inputs_embeds=inputs_embeds,
pixel_values=pixel_values,
image_sizes=image_sizes,
attention_mask=attention_mask,
**kwargs,
)
else:
use_flash_attention = kwargs.get("use_flash_attention", False)
flash_attention_recompute = kwargs.get("flash_attention_recompute", False)
position_ids = kwargs.get("position_ids", None)
labels = kwargs.get("labels", None)
if past_key_values is None and pixel_values is not None and input_ids.shape[1] != 1:
vision_feature_select_strategy = kwargs.get("vision_feature_select_strategy", None)
vision_feature_layer = kwargs.get("vision_feature_layer", None)
vision_feature_select_strategy = (
vision_feature_select_strategy
if vision_feature_select_strategy is not None
else self.config.vision_feature_select_strategy
)
vision_feature_layer = (
vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer
)
# 1. Extract the input embeddings # 1. Extract the input embeddings
inputs_embeds = self.get_input_embeddings()(input_ids)
# 2. Merge text and images # 2. Merge text and images
batch_size, num_patches, num_channels, height, width = pixel_values.shape num_images, num_patches, channels, height, width = pixel_values.shape
reshaped_pixel_values = pixel_values.view(batch_size * num_patches, num_channels, height, width) pixel_values = pixel_values.view(
image_features = self.vision_tower( num_images * num_patches, channels, height, width
reshaped_pixel_values,
output_hidden_states=True,
use_flash_attention=use_flash_attention,
flash_attention_recompute=flash_attention_recompute,
) )
image_features = self.vision_tower(pixel_values)
selected_image_feature = image_features.hidden_states[vision_feature_layer] # selected_image_feature = image_features.hidden_states[self.config.vision_feature_layer]
# Already done within the clip model
selected_image_feature = image_features.last_hidden_state
if vision_feature_select_strategy == "default": if self.config.vision_feature_select_strategy == "default":
selected_image_feature = selected_image_feature[:, 1:] selected_image_feature = selected_image_feature[:, 1:]
elif vision_feature_select_strategy == "full": elif self.config.vision_feature_select_strategy == "full":
selected_image_feature = selected_image_feature selected_image_feature = selected_image_feature
else:
raise RuntimeError(
f"Strategy `{self.config.vision_feature_select_strategy}` is not supported/valid."
)
image_features = self.multi_modal_projector(selected_image_feature) image_features = self.multi_modal_projector(selected_image_feature)
# split up image_features for each of the individual images # split up image_features for each of the individual images
# hence we get a list of image_features, each of shape (5, num_patches, hidden_size) # hence we get a list of image_features, each of shape (5, num_patches, hidden_size)
# if we assume each image has 5 image features (base image + 4 patches) # if we assume each image has 5 image features (base image + 4 patches)
split_sizes = [image.shape[0] for image in pixel_values] split_sizes = [num_patches] * num_images
image_features = torch.split(image_features, split_sizes, dim=0) image_features = torch.split(image_features, split_sizes, dim=0)
# NOTE we only support multimodal_patch_merge_type == "spatial_unpad" # NOTE we only support multimodal_patch_merge_type == "spatial_unpad"
height = width = self.config.vision_config.image_size // self.config.vision_config.patch_size height = width = (
self.config.vision_config.image_size
// self.config.vision_config.patch_size
)
new_image_features = [] new_image_features = []
for image_idx, image_feature in enumerate(image_features): for image_idx, image_feature in enumerate(image_features):
@ -203,94 +226,58 @@ class LlavaNextForConditionalGeneration(GaudiLlavaNextForConditionalGeneration):
image_feature = image_feature[1:] image_feature = image_feature[1:]
if height * width != base_image_feature.shape[0]: if height * width != base_image_feature.shape[0]:
raise ValueError("The number of patches is not consistent with the image size.") raise ValueError(
"The number of patches is not consistent with the image size."
)
num_patch_height, num_patch_width = get_anyres_image_grid_shape( num_patch_height, num_patch_width = get_anyres_image_grid_shape(
image_sizes[image_idx].tolist(), image_sizes[image_idx],
self.config.image_grid_pinpoints, self.config.image_grid_pinpoints,
self.config.vision_config.image_size, self.config.vision_config.image_size,
) )
image_feature = image_feature.view(
image_feature = image_feature.view(num_patch_height, num_patch_width, height, width, -1) num_patch_height, num_patch_width, height, width, -1
)
image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous() image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous()
image_feature = image_feature.flatten(1, 2).flatten(2, 3) image_feature = image_feature.flatten(1, 2).flatten(2, 3)
image_feature = unpad_image(image_feature, image_sizes[image_idx]) image_feature = unpad_image(image_feature, image_sizes[image_idx])
image_feature = torch.cat( image_feature = torch.cat(
( (
image_feature, image_feature,
self.image_newline[:, None, None].expand(*image_feature.shape[:-1], 1), self.image_newline[:, None, None].expand(
*image_feature.shape[:-1], 1
),
), ),
dim=-1, dim=-1,
) )
image_feature = image_feature.flatten(1, 2).transpose(0, 1) image_feature = image_feature.flatten(1, 2).transpose(0, 1)
image_feature = torch.cat((base_image_feature, image_feature), dim=0) image_feature = torch.cat(
(base_image_feature, image_feature), dim=0
)
else: else:
image_feature = image_feature[0] image_feature = image_feature[0]
image_feature = torch.cat((image_feature, self.image_newline[None]), dim=0) image_feature = torch.cat(
(image_feature, self.image_newline[None]), dim=0
)
new_image_features.append(image_feature) new_image_features.append(image_feature)
image_features = torch.stack(new_image_features, dim=0) image_features = torch.stack(new_image_features, dim=0)
inputs_embeds = self._merge_input_ids_with_image_features(inputs_embeds, image_features, input_ids)
self.image_offset = image_features.shape[1] - 1 # image_token has occupied 1 token position.
# In case input_ids.shape[1] == 1 & pixel_values==None & past_key_values != None, we are in the case of
# generation with cache
elif past_key_values is not None:
seq_len = input_ids.shape[1]
pad_len = seq_len - token_idx.item()
input_ids = torch.index_select(input_ids, 1, token_idx - 1)
# Retrieve the first layer to inspect the logits and mask out the hidden states
# that are set to 0
first_layer_past_key_value = past_key_values[0][0][:, :, :, 0]
# Sum all dimensions of head_dim (-2) to avoid random errors such as: https://github.com/huggingface/transformers/pull/28032#issuecomment-1863691941 inputs_embeds = self._merge_input_ids_with_image_features(
batch_index, non_attended_tokens = torch.where(first_layer_past_key_value.float().sum(-2) == 0) input_ids, inputs_embeds, image_features
# Get the target length
past_length = first_layer_past_key_value.shape[-1]
extended_attention_mask = torch.ones(
(attention_mask.shape[0], past_length),
dtype=attention_mask.dtype,
device=attention_mask.device,
)
# Filter out only the tokens that can be un-attended, this can happen
# if one uses Llava + Fused modules where the cache on the
# first iteration is already big enough, or if one passes custom cache
valid_indices = non_attended_tokens < extended_attention_mask.size(-1)
new_batch_index = batch_index[valid_indices]
new_non_attended_tokens = non_attended_tokens[valid_indices]
# Zero-out the places where we don't need to attend
extended_attention_mask[new_batch_index, new_non_attended_tokens] = 0
attention_mask = extended_attention_mask
attention_mask[:, -pad_len:] = 0
if attention_mask is not None and position_ids is None:
# create position_ids on the fly for batch generation
position_ids = attention_mask.long().cumsum(-1) - 1
position_ids.masked_fill_(attention_mask == 0, 1)
if past_key_values:
if token_idx is not None:
position_ids = torch.sum(attention_mask, dim=1).unsqueeze(-1) - 1
else:
position_ids = position_ids[:, -input_ids.shape[1] :]
# if `inputs_embeds` are passed, we only want to use them in the 1st generation step
if inputs_embeds is not None and past_key_values is None:
model_inputs = {"inputs_embeds": inputs_embeds}
else:
model_inputs = {"input_ids": input_ids}
model_inputs.update(
{
"position_ids": position_ids,
"past_key_values": past_key_values,
"use_cache": kwargs.get("use_cache"),
"attention_mask": attention_mask,
"token_idx": token_idx,
"labels": labels,
"use_flash_attention": use_flash_attention,
"flash_attention_recompute": flash_attention_recompute,
}
) )
return model_inputs hidden_states = self.language_model.model(
inputs_embeds=inputs_embeds,
position_ids=position_ids,
cu_seqlen_prefill=cu_seqlen_prefill,
kv_cache=kv_cache,
block_tables=block_tables,
slots=slots,
input_lengths=input_lengths,
max_s=max_s,
true_max_s=max_s,
prefill_cache_indices=None,
)
if lm_head_indices is not None:
hidden_states = hidden_states[lm_head_indices]
logits, speculative_logits = self.language_model.lm_head(hidden_states)
return logits, speculative_logits

2
server/text_generation_server/models/flash_mistral.py
View File

@ -10,7 +10,7 @@ from transformers import PreTrainedTokenizerBase, AutoTokenizer, AutoConfig
from typing import Optional, Tuple, Type from typing import Optional, Tuple, Type
from text_generation_server.pb import generate_pb2 from text_generation_server.pb import generate_pb2
from text_generation_server.models.flash_causal_lm import FlashCausalLM from text_generation_server.models import FlashCausalLM
from text_generation_server.models.flash_causal_lm import FlashCausalLMBatch, BLOCK_SIZE from text_generation_server.models.flash_causal_lm import FlashCausalLMBatch, BLOCK_SIZE
from text_generation_server.models.cache_manager import ( from text_generation_server.models.cache_manager import (
get_cache_manager, get_cache_manager,

28
server/text_generation_server/models/model.py
View File

@ -2,8 +2,8 @@ import inspect
import torch import torch
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from typing import List, Optional, Tuple, Type, TypeVar from typing import List, Tuple, Optional, TypeVar, Type
from transformers import PreTrainedTokenizerBase from transformers import PreTrainedTokenizerBase, PretrainedConfig
from text_generation_server.models.types import Batch, Generation from text_generation_server.models.types import Batch, Generation
from text_generation_server.utils.speculate import get_speculate from text_generation_server.utils.speculate import get_speculate
@ -22,10 +22,10 @@ class Model(ABC):
device: torch.device, device: torch.device,
rank: int = 0, rank: int = 0,
world_size: int = 1, world_size: int = 1,
kwargs: dict = {}, sliding_window: Optional[int] = None,
speculate: Optional[int] = None, speculate: Optional[int] = None,
): ):
self.model = model self.model = model.eval()
self.tokenizer = tokenizer self.tokenizer = tokenizer
# all_special_ids is not set correctly if the rust tokenizer is unpacked # all_special_ids is not set correctly if the rust tokenizer is unpacked
@ -40,7 +40,8 @@ class Model(ABC):
self.device = device self.device = device
self.rank = rank self.rank = rank
self.world_size = world_size self.world_size = world_size
self.kwargs = kwargs self.sliding_window = sliding_window if sliding_window != -1 else None
if speculate is None: if speculate is None:
speculate = get_speculate() speculate = get_speculate()
self.speculate = speculate self.speculate = speculate
@ -54,10 +55,14 @@ class Model(ABC):
@property @property
def info(self) -> InfoResponse: def info(self) -> InfoResponse:
if self.requires_padding and self.sliding_window is not None:
raise NotImplementedError("sliding_window is not implemented with padding")
return InfoResponse( return InfoResponse(
requires_padding=self.requires_padding, requires_padding=self.requires_padding,
dtype=str(self.dtype), dtype=str(self.dtype),
device_type=self.device.type, device_type=self.device.type,
window_size=self.sliding_window,
speculate=self.speculate, speculate=self.speculate,
) )
@ -72,21 +77,28 @@ class Model(ABC):
) -> Tuple[List[Generation], Optional[B], Tuple[int, int]]: ) -> Tuple[List[Generation], Optional[B], Tuple[int, int]]:
raise NotImplementedError raise NotImplementedError
def warmup(self, batch: B, max_total_tokens: int): def warmup(self, batch: B) -> Optional[int]:
self.generate_token(batch) self.generate_token(batch)
return None
def decode_token( def decode_token(
self, self,
all_input_ids: List[int], all_input_ids: List[int],
prefix_offset: int = 0, prefix_offset: int = 0,
read_offset: int = 0, read_offset: int = 0,
skip_special_tokens: bool = False,
) -> Tuple[str, int, int]: ) -> Tuple[str, int, int]:
"""Hack to hopefully support generate_stream for the maximum number of tokenizers""" """Hack to hopefully support generate_stream for the maximum number of tokenizers"""
# The prefix text is necessary only to defeat cleanup algorithms in the decode # The prefix text is necessary only to defeat cleanup algorithms in the decode
# which decide to add a space or not depending on the surrounding ids. # which decide to add a space or not depending on the surrounding ids.
prefix_text = self.tokenizer.decode(all_input_ids[prefix_offset:read_offset], skip_special_tokens=False) prefix_text = self.tokenizer.decode(
new_text = self.tokenizer.decode(all_input_ids[prefix_offset:], skip_special_tokens=False) all_input_ids[prefix_offset:read_offset],
skip_special_tokens=skip_special_tokens,
)
new_text = self.tokenizer.decode(
all_input_ids[prefix_offset:], skip_special_tokens=skip_special_tokens
)
if len(new_text) > len(prefix_text) and not new_text.endswith("<EFBFBD>"): if len(new_text) > len(prefix_text) and not new_text.endswith("<EFBFBD>"):
# utf-8 char at the end means it's a potential unfinished byte sequence # utf-8 char at the end means it's a potential unfinished byte sequence

47
server/text_generation_server/models/santacoder.py
View File

@ -1,5 +1,8 @@
from typing import Optional, List
import torch import torch
import torch.distributed
from typing import Optional, List
from transformers import AutoTokenizer, AutoModelForCausalLM
from text_generation_server.models import CausalLM from text_generation_server.models import CausalLM
@ -15,19 +18,29 @@ class SantaCoder(CausalLM):
self, self,
model_id: str, model_id: str,
revision: Optional[str] = None, revision: Optional[str] = None,
quantize: Optional[str] = None,
speculator: Optional[str] = None, speculator: Optional[str] = None,
dtype: Optional[torch.dtype] = None, dtype: Optional[torch.dtype] = None,
trust_remote_code: bool = False, trust_remote_code: bool = False,
): ):
super().__init__( if torch.cuda.is_available():
model_id=model_id, device = torch.device("cuda")
dtype = torch.float16 if dtype is None else dtype
else:
if quantize:
raise ValueError("quantization is not available on CPU")
device = torch.device("cpu")
dtype = torch.float32 if dtype is None else dtype
tokenizer = AutoTokenizer.from_pretrained(
model_id,
revision=revision, revision=revision,
speculator=speculator, padding_side="left",
dtype=dtype, truncation_side="left",
trust_remote_code=trust_remote_code, trust_remote_code=trust_remote_code,
) )
tokenizer.add_special_tokens(
self.tokenizer.add_special_tokens(
{ {
"additional_special_tokens": [ "additional_special_tokens": [
EOD, EOD,
@ -39,7 +52,25 @@ class SantaCoder(CausalLM):
"pad_token": EOD, "pad_token": EOD,
} }
) )
with device:
model = AutoModelForCausalLM.from_pretrained(
model_id,
revision=revision,
torch_dtype=dtype,
load_in_8bit=quantize == "bitsandbytes",
trust_remote_code=trust_remote_code,
)
super(CausalLM, self).__init__(
model=model,
tokenizer=tokenizer,
requires_padding=True,
dtype=dtype,
device=device,
)
def decode(self, generated_ids: List[int]) -> str: def decode(self, generated_ids: List[int]) -> str:
# Do not skip special tokens as they are used for custom parsing rules of the generated text # Do not skip special tokens as they are used for custom parsing rules of the generated text
return self.tokenizer.decode(generated_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False) return self.tokenizer.decode(
generated_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False
)

47
server/text_generation_server/models/starcoder.py
View File

@ -1,47 +0,0 @@
from loguru import logger
import torch
from dataclasses import dataclass
import os
from typing import List, Optional, Type
from text_generation_server.models import CausalLM
from text_generation_server.models.causal_lm import CausalLMBatch
@dataclass
class StarCoderCausalLMBatch(CausalLMBatch):
past_key_values: Optional[List[torch.Tensor]]
def detach_kv_cache(self):
past_keys = []
past_values = []
last_dim = int(self.past_key_values[0].size(dim=-1)/2)
for key_value in self.past_key_values:
past_keys.append(key_value.split((last_dim, last_dim), dim=-1)[0])
past_values.append(key_value.split((last_dim, last_dim), dim=-1)[1])
del self.past_key_values
return past_keys, past_values
def attach_kv_cache(self, past_keys, past_values):
self.past_key_values = [
torch.cat((key, value), dim=-1) for key, value in zip(past_keys, past_values)]
class StarCoder(CausalLM):
def __init__(
self,
model_id: str,
revision: Optional[str] = None,
dtype: Optional[torch.dtype] = None,
):
super(StarCoder, self).__init__(
model_id=model_id,
revision=revision,
dtype=dtype,
)
@property
def batch_type(self) -> Type[CausalLMBatch]:
return StarCoderCausalLMBatch

1158
server/text_generation_server/models/vlm_causal_lm.py
View File

File diff suppressed because it is too large Load Diff

121
server/text_generation_server/server.py
View File

@ -1,8 +1,5 @@
# Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.
import asyncio import asyncio
import os import os
import sys
import torch import torch
import time import time
import signal import signal
@ -17,9 +14,6 @@ from typing import List, Optional
from text_generation_server.cache import Cache from text_generation_server.cache import Cache
from text_generation_server.interceptor import ExceptionInterceptor from text_generation_server.interceptor import ExceptionInterceptor
from text_generation_server.models import Model, get_model from text_generation_server.models import Model, get_model
from text_generation_server.pb import generate_pb2_grpc, generate_pb2
from text_generation_server.tracing import UDSOpenTelemetryAioServerInterceptor
from text_generation_server.models.globals import set_model_id
try: try:
from text_generation_server.models.pali_gemma import PaliGemmaBatch from text_generation_server.models.pali_gemma import PaliGemmaBatch
@ -32,7 +26,10 @@ try:
except (ImportError, NotImplementedError): except (ImportError, NotImplementedError):
# These imports can fail on CPU/Non flash. # These imports can fail on CPU/Non flash.
VLM_BATCH_TYPES = set() VLM_BATCH_TYPES = set()
from text_generation_server.utils.version import is_driver_compatible, MIN_TGI_GAUDI_SYNAPSE_VERSION
from text_generation_server.pb import generate_pb2_grpc, generate_pb2
from text_generation_server.tracing import UDSOpenTelemetryAioServerInterceptor
from text_generation_server.models.globals import set_model_id
class SignalHandler: class SignalHandler:
@ -52,25 +49,24 @@ class TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer):
self, self,
model: Model, model: Model,
cache: Cache, cache: Cache,
quantize: Optional[str],
server_urls: List[str], server_urls: List[str],
): ):
self.cache = cache self.cache = cache
self.model = model self.model = model
self.quantize = quantize
self.server_urls = server_urls self.server_urls = server_urls
# For some reason, inference_mode does not work well with GLOO which we use on CPU # For some reason, inference_mode does not work well with GLOO which we use on CPU
# TODO: The inferecemode set messes up the autograd op dispatch. And results in aten::matmul if model.device.type == "cuda":
# op not optimized issue. Will investigate further.
# if model.device.type == "hpu":
# Force inference mode for the lifetime of TextGenerationService # Force inference mode for the lifetime of TextGenerationService
# self._inference_mode_raii_guard = torch._C._InferenceMode(True) self._inference_mode_raii_guard = torch._C._InferenceMode(True)
async def Info(self, request, context): async def Info(self, request, context):
return self.model.info return self.model.info
async def Health(self, request, context): async def Health(self, request, context):
if self.model.device.type == "hpu": if self.model.device.type == "cuda":
torch.zeros((2, 2)).to("hpu") torch.zeros((2, 2)).cuda()
return generate_pb2.HealthResponse() return generate_pb2.HealthResponse()
async def ServiceDiscovery(self, request, context): async def ServiceDiscovery(self, request, context):
@ -93,19 +89,41 @@ class TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer):
return generate_pb2.FilterBatchResponse(batch=filtered_batch.to_pb()) return generate_pb2.FilterBatchResponse(batch=filtered_batch.to_pb())
async def Warmup(self, request, context): async def Warmup(self, request, context):
def batch_from_pb(batch): if self.quantize == "gptq":
return self.model.batch_type.from_pb( try:
batch, self.model.tokenizer, self.model.dtype, self.model.device # When using GPTQ, Exllama kernels need some global kernels
# For which we have the finale shapes only after the model has loaded
# This will allocate those buffers.
from text_generation_server.layers.gptq import (
create_exllama_buffers,
set_device,
) )
if self.model.batch_type in VLM_BATCH_TYPES : set_device(self.model.device)
max_supported_total_tokens = self.model.warmup(request) create_exllama_buffers(request.max_prefill_tokens)
return generate_pb2.WarmupResponse(max_supported_total_tokens=max_supported_total_tokens) except ImportError:
else: pass
batches = [batch_from_pb(batch) for batch in request.batches]
self.model.warmup(batches)
return generate_pb2.WarmupResponse()
if (
self.model.batch_type in VLM_BATCH_TYPES
): # Hack, i would rather use kwargs in the `from_pb` call
batch = self.model.batch_type.from_pb_processor(
request.batch,
self.model.tokenizer,
self.model.processor,
self.model.model.config,
self.model.dtype,
self.model.device,
)
else:
batch = self.model.batch_type.from_pb(
request.batch, self.model.tokenizer, self.model.dtype, self.model.device
)
max_supported_total_tokens = self.model.warmup(batch)
return generate_pb2.WarmupResponse(
max_supported_total_tokens=max_supported_total_tokens
)
async def Prefill(self, request, context): async def Prefill(self, request, context):
start = time.time_ns() start = time.time_ns()
@ -125,7 +143,7 @@ class TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer):
request.batch, self.model.tokenizer, self.model.dtype, self.model.device request.batch, self.model.tokenizer, self.model.dtype, self.model.device
) )
generations, next_batch, timings = self.model.generate_token([batch]) generations, next_batch, timings = self.model.generate_token(batch)
self.cache.set(next_batch) self.cache.set(next_batch)
return generate_pb2.PrefillResponse( return generate_pb2.PrefillResponse(
@ -151,13 +169,21 @@ class TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer):
if len(batches) == 0: if len(batches) == 0:
raise ValueError("All batches are empty") raise ValueError("All batches are empty")
generations, next_batch, timings = self.model.generate_token(batches) if len(batches) > 1:
start_concat = time.time_ns()
batch = self.model.batch_type.concatenate(batches)
concat_ns = time.time_ns() - start_concat
else:
batch = batches[0]
concat_ns = None
generations, next_batch, timings = self.model.generate_token(batch)
self.cache.set(next_batch) self.cache.set(next_batch)
return generate_pb2.DecodeResponse( return generate_pb2.DecodeResponse(
generations=[generation.to_pb() for generation in generations], generations=[generation.to_pb() for generation in generations],
batch=next_batch.to_pb() if next_batch else None, batch=next_batch.to_pb() if next_batch else None,
concat_ns=None, # TODO: measure concat time concat_ns=concat_ns,
forward_ns=timings[0], forward_ns=timings[0],
decode_ns=timings[1], decode_ns=timings[1],
total_ns=time.time_ns() - start, total_ns=time.time_ns() - start,
@ -168,62 +194,41 @@ def serve(
model_id: str, model_id: str,
revision: Optional[str], revision: Optional[str],
sharded: bool, sharded: bool,
quantize: Optional[str],
speculate: Optional[int], speculate: Optional[int],
dtype: Optional[str], dtype: Optional[str],
trust_remote_code: bool, trust_remote_code: bool,
uds_path: Path, uds_path: Path,
): ):
# Remove default handler
#logger.remove()
logger.add(
sys.stdout,
format="{message}",
filter="text_generation_server",
level="INFO",
serialize=False,
backtrace=True,
diagnose=False,
)
async def serve_inner( async def serve_inner(
model_id: str, model_id: str,
revision: Optional[str], revision: Optional[str],
sharded: bool = False, sharded: bool = False,
quantize: Optional[str] = None,
speculate: Optional[int] = None, speculate: Optional[int] = None,
dtype: Optional[str] = None, dtype: Optional[str] = None,
trust_remote_code: bool = False, trust_remote_code: bool = False,
): ):
if not is_driver_compatible():
logger.warning(f"Current Synapse version is lower than the minimum version supported: {MIN_TGI_GAUDI_SYNAPSE_VERSION}, this could result in failures")
unix_socket_template = "unix://{}-{}" unix_socket_template = "unix://{}-{}"
logger.info("Server:server_inner: sharded ={}".format(sharded))
if sharded: if sharded:
rank = int(os.environ["RANK"])
logger.info("Server:server_inner: rank ={}".format(rank))
server_urls = [ server_urls = [
unix_socket_template.format(uds_path, rank) for rank in range(int(os.environ["WORLD_SIZE"])) unix_socket_template.format(uds_path, rank)
for rank in range(int(os.environ["WORLD_SIZE"]))
] ]
local_url = server_urls[int(os.environ["RANK"])] local_url = server_urls[int(os.environ["RANK"])]
else: else:
local_url = unix_socket_template.format(uds_path, 0) local_url = unix_socket_template.format(uds_path, 0)
server_urls = [local_url] server_urls = [local_url]
logger.info("Server:server_inner: data type = {}, local_url = {}".format(dtype, local_url))
if dtype == "bfloat16" or None:
data_type = torch.bfloat16
else:
data_type = torch.float
if revision == "None":
revision = None
try: try:
model = get_model( model = get_model(
model_id, model_id,
revision, revision,
sharded,
quantize,
speculate, speculate,
data_type, dtype,
trust_remote_code trust_remote_code,
) )
except Exception: except Exception:
logger.exception("Error when initializing model") logger.exception("Error when initializing model")
@ -236,7 +241,7 @@ def serve(
] ]
) )
generate_pb2_grpc.add_TextGenerationServiceServicer_to_server( generate_pb2_grpc.add_TextGenerationServiceServicer_to_server(
TextGenerationService(model, Cache(), server_urls), server TextGenerationService(model, Cache(), quantize, server_urls), server
) )
SERVICE_NAMES = ( SERVICE_NAMES = (
generate_pb2.DESCRIPTOR.services_by_name["TextGenerationService"].full_name, generate_pb2.DESCRIPTOR.services_by_name["TextGenerationService"].full_name,
@ -255,6 +260,6 @@ def serve(
set_model_id(model_id) set_model_id(model_id)
asyncio.run( asyncio.run(
serve_inner( serve_inner(
model_id, revision, sharded, speculate, dtype, trust_remote_code model_id, revision, sharded, quantize, speculate, dtype, trust_remote_code
) )
) )

37
server/text_generation_server/tgi_service.py
View File

@ -1,37 +0,0 @@
import os
from pathlib import Path
from loguru import logger
import sys
from text_generation_server import server
import argparse
def main(args):
logger.info("TGIService: starting tgi service .... ")
logger.info(
"TGIService: --model_id {}, --revision {}, --sharded {}, --speculate {}, --dtype {}, --trust_remote_code {}, --uds_path {} ".format(
args.model_id, args.revision, args.sharded, args.speculate, args.dtype, args.trust_remote_code, args.uds_path
)
)
server.serve(
model_id=args.model_id,
revision=args.revision,
sharded=args.sharded,
speculate=args.speculate,
dtype=args.dtype,
trust_remote_code=args.trust_remote_code,
uds_path=args.uds_path,
)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--model_id", type=str)
parser.add_argument("--revision", type=str)
parser.add_argument("--sharded", type=bool)
parser.add_argument("--speculate", type=int, default=None)
parser.add_argument("--dtype", type=str)
parser.add_argument("--trust_remote_code", type=bool)
parser.add_argument("--uds_path", type=Path)
args = parser.parse_args()
main(args)

6
server/text_generation_server/utils/__init__.py
View File

@ -1,6 +1,3 @@
# Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.
import text_generation_server.habana_quantization_env
from text_generation_server.utils.convert import convert_file, convert_files from text_generation_server.utils.convert import convert_file, convert_files
from text_generation_server.utils.dist import initialize_torch_distributed from text_generation_server.utils.dist import initialize_torch_distributed
from text_generation_server.utils.weights import Weights from text_generation_server.utils.weights import Weights
@ -21,9 +18,6 @@ from text_generation_server.utils.tokens import (
FinishReason, FinishReason,
Sampling, Sampling,
Greedy, Greedy,
make_tokenizer_optional,
is_tokenizer_transparent,
pad_next_token_chooser_parameters,
) )
__all__ = [ __all__ = [

31
server/text_generation_server/utils/debug.py
View File

@ -1,31 +0,0 @@
# Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.
import os
import glob
import time
from optimum.habana.utils import to_gb_rounded
import habana_frameworks.torch as htorch
START_TS = None
DBG_TRACE_FILENAME = os.environ.get('DBG_TRACE_FILENAME')
if 'GRAPH_VISUALIZATION' in os.environ:
for f in glob.glob('.graph_dumps/*'):
os.remove(f)
def count_hpu_graphs():
return len(glob.glob('.graph_dumps/*PreGraph*'))
def dbg_trace(tag, txt):
global START_TS
if DBG_TRACE_FILENAME is not None and int(os.getenv("RANK", 0)) == 0:
if START_TS is None:
START_TS = time.perf_counter()
time_offset = time.perf_counter() - START_TS
mem_stats = htorch.hpu.memory.memory_stats()
mem_used = to_gb_rounded(mem_stats['InUse'])
max_mem_used = to_gb_rounded(mem_stats['MaxInUse'])
print(f'ts:{time_offset:.3f}s g:{count_hpu_graphs()} mu:{mem_used:.1f}GB '
f'mmu:{max_mem_used:.1f}GB | {tag} | {txt}', flush=True, file=open(DBG_TRACE_FILENAME, 'a'))

11
server/text_generation_server/utils/dist.py
View File

@ -44,12 +44,6 @@ class FakeGroup:
def initialize_torch_distributed(): def initialize_torch_distributed():
import habana_frameworks.torch.core as htcore
rank = int(os.getenv("RANK", "0"))
world_size = int(os.getenv("WORLD_SIZE", "1"))
options = None
if torch.cuda.is_available(): if torch.cuda.is_available():
from torch.distributed import ProcessGroupNCCL from torch.distributed import ProcessGroupNCCL
@ -62,11 +56,6 @@ def initialize_torch_distributed():
options = ProcessGroupNCCL.Options() options = ProcessGroupNCCL.Options()
options.is_high_priority_stream = True options.is_high_priority_stream = True
options._timeout = timedelta(seconds=60) options._timeout = timedelta(seconds=60)
elif torch.hpu.is_available():
backend = "hccl"
n_hpus = torch.hpu.device_count()
if world_size > n_hpus:
raise ValueError(f"WORLD_SIZE ({world_size}) is higher than the number of available HPUs ({n_hpus}).")
else: else:
try: try:
import oneccl_bindings_for_pytorch import oneccl_bindings_for_pytorch

359
server/text_generation_server/utils/gptq/quant_linear.py
View File

@ -1,359 +0,0 @@
import math
import numpy as np
import torch
import torch.nn as nn
from torch.cuda.amp import custom_bwd, custom_fwd
try:
import triton
import triton.language as tl
from . import custom_autotune
# code based https://github.com/fpgaminer/GPTQ-triton
@custom_autotune.autotune(
configs=[
triton.Config(
{
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 32,
"GROUP_SIZE_M": 8,
},
num_stages=4,
num_warps=4,
),
triton.Config(
{
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 32,
"GROUP_SIZE_M": 8,
},
num_stages=4,
num_warps=4,
),
triton.Config(
{
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 32,
"GROUP_SIZE_M": 8,
},
num_stages=4,
num_warps=4,
),
triton.Config(
{
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 32,
"GROUP_SIZE_M": 8,
},
num_stages=4,
num_warps=4,
),
triton.Config(
{
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 32,
"GROUP_SIZE_M": 8,
},
num_stages=4,
num_warps=4,
),
triton.Config(
{
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 32,
"GROUP_SIZE_M": 8,
},
num_stages=2,
num_warps=8,
),
triton.Config(
{
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 8,
},
num_stages=3,
num_warps=8,
),
triton.Config(
{
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 8,
},
num_stages=2,
num_warps=4,
),
],
key=["M", "N", "K"],
nearest_power_of_two=True,
prune_configs_by={
"early_config_prune": custom_autotune.matmul248_kernel_config_pruner,
"perf_model": None,
"top_k": None,
},
)
@triton.jit
def matmul_248_kernel(
a_ptr,
b_ptr,
c_ptr,
scales_ptr,
zeros_ptr,
g_ptr,
M,
N,
K,
bits,
maxq,
stride_am,
stride_ak,
stride_bk,
stride_bn,
stride_cm,
stride_cn,
stride_scales,
stride_zeros,
BLOCK_SIZE_M: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr,
GROUP_SIZE_M: tl.constexpr,
):
"""
Compute the matrix multiplication C = A x B.
A is of shape (M, K) float16
B is of shape (K//8, N) int32
C is of shape (M, N) float16
scales is of shape (G, N) float16
zeros is of shape (G, N) float16
g_ptr is of shape (K) int32
"""
infearure_per_bits = 32 // bits
pid = tl.program_id(axis=0)
num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
num_pid_k = tl.cdiv(K, BLOCK_SIZE_K)
num_pid_in_group = GROUP_SIZE_M * num_pid_n
group_id = pid // num_pid_in_group
first_pid_m = group_id * GROUP_SIZE_M
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
pid_m = first_pid_m + (pid % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m
offs_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
offs_k = tl.arange(0, BLOCK_SIZE_K)
a_ptrs = a_ptr + (
offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak
) # (BLOCK_SIZE_M, BLOCK_SIZE_K)
a_mask = offs_am[:, None] < M
# b_ptrs is set up such that it repeats elements along the K axis 8 times
b_ptrs = b_ptr + (
(offs_k[:, None] // infearure_per_bits) * stride_bk
+ offs_bn[None, :] * stride_bn
) # (BLOCK_SIZE_K, BLOCK_SIZE_N)
g_ptrs = g_ptr + offs_k
# shifter is used to extract the N bits of each element in the 32-bit word from B
scales_ptrs = scales_ptr + offs_bn[None, :]
zeros_ptrs = zeros_ptr + (offs_bn[None, :] // infearure_per_bits)
shifter = (offs_k % infearure_per_bits) * bits
zeros_shifter = (offs_bn % infearure_per_bits) * bits
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
for k in range(0, num_pid_k):
g_idx = tl.load(g_ptrs)
# Fetch scales and zeros; these are per-outfeature and thus reused in the inner loop
scales = tl.load(
scales_ptrs + g_idx[:, None] * stride_scales
) # (BLOCK_SIZE_K, BLOCK_SIZE_N,)
zeros = tl.load(
zeros_ptrs + g_idx[:, None] * stride_zeros
) # (BLOCK_SIZE_K, BLOCK_SIZE_N,)
zeros = (zeros >> zeros_shifter[None, :]) & maxq
zeros = (zeros + 1) & maxq # eventually avoid overflow
a = tl.load(a_ptrs, mask=a_mask, other=0.0) # (BLOCK_SIZE_M, BLOCK_SIZE_K)
b = tl.load(b_ptrs) # (BLOCK_SIZE_K, BLOCK_SIZE_N), but repeated
# Now we need to unpack b (which is N-bit values) into 32-bit values
b = (b >> shifter[:, None]) & maxq # Extract the N-bit values
b = (b - zeros) * scales # Scale and shift
accumulator += tl.dot(a, b)
a_ptrs += BLOCK_SIZE_K
b_ptrs += (BLOCK_SIZE_K // infearure_per_bits) * stride_bk
g_ptrs += BLOCK_SIZE_K
c_ptrs = c_ptr + stride_cm * offs_am[:, None] + stride_cn * offs_bn[None, :]
c_mask = (offs_am[:, None] < M) & (offs_bn[None, :] < N)
tl.store(c_ptrs, accumulator, mask=c_mask)
except:
print("triton not installed.")
def matmul248(input, qweight, scales, qzeros, g_idx, bits, maxq):
with torch.cuda.device(input.device):
output = torch.empty(
(input.shape[0], qweight.shape[1]), device=input.device, dtype=torch.float16
)
grid = lambda META: (
triton.cdiv(input.shape[0], META["BLOCK_SIZE_M"])
* triton.cdiv(qweight.shape[1], META["BLOCK_SIZE_N"]),
)
matmul_248_kernel[grid](
input,
qweight,
output,
scales,
qzeros,
g_idx,
input.shape[0],
qweight.shape[1],
input.shape[1],
bits,
maxq,
input.stride(0),
input.stride(1),
qweight.stride(0),
qweight.stride(1),
output.stride(0),
output.stride(1),
scales.stride(0),
qzeros.stride(0),
)
return output
class QuantLinearFunction(torch.autograd.Function):
@staticmethod
@custom_fwd(cast_inputs=torch.float16)
def forward(ctx, input, qweight, scales, qzeros, g_idx, bits, maxq):
output = matmul248(input, qweight, scales, qzeros, g_idx, bits, maxq)
return output
class QuantLinear(nn.Module):
def __init__(self, qweight, qzeros, scales, g_idx, bias, bits, groupsize):
super().__init__()
self.register_buffer("qweight", qweight)
self.register_buffer("qzeros", qzeros)
self.register_buffer("scales", scales)
self.register_buffer("g_idx", g_idx)
if bias is not None:
self.register_buffer("bias", bias)
else:
self.bias = None
if bits not in [2, 4, 8]:
raise NotImplementedError("Only 2,4,8 bits are supported.")
self.bits = bits
self.maxq = 2**self.bits - 1
self.groupsize = groupsize
self.outfeatures = qweight.shape[1]
self.infeatures = qweight.shape[0] * 32 // bits
@classmethod
def new(cls, bits, groupsize, infeatures, outfeatures, bias):
if bits not in [2, 4, 8]:
raise NotImplementedError("Only 2,4,8 bits are supported.")
qweight = torch.zeros((infeatures // 32 * bits, outfeatures), dtype=torch.int32)
qzeros = torch.zeros(
(math.ceil(infeatures / groupsize), outfeatures // 32 * bits),
dtype=torch.int32,
)
scales = torch.zeros(
(math.ceil(infeatures / groupsize), outfeatures), dtype=torch.float16
)
g_idx = torch.tensor(
[i // groupsize for i in range(infeatures)], dtype=torch.int32
)
if bias:
bias = torch.zeros((outfeatures), dtype=torch.float16)
else:
bias = None
return cls(qweight, qzeros, scales, g_idx, bias, bits, groupsize)
def pack(self, linear, scales, zeros, g_idx=None):
self.g_idx = g_idx.clone() if g_idx is not None else self.g_idx
scales = scales.t().contiguous()
zeros = zeros.t().contiguous()
scale_zeros = zeros * scales
self.scales = scales.clone().half()
if linear.bias is not None:
self.bias = linear.bias.clone().half()
intweight = []
for idx in range(self.infeatures):
intweight.append(
torch.round(
(linear.weight.data[:, idx] + scale_zeros[self.g_idx[idx]])
/ self.scales[self.g_idx[idx]]
).to(torch.int)[:, None]
)
intweight = torch.cat(intweight, dim=1)
intweight = intweight.t().contiguous()
intweight = intweight.numpy().astype(np.uint32)
qweight = np.zeros(
(intweight.shape[0] // 32 * self.bits, intweight.shape[1]), dtype=np.uint32
)
i = 0
row = 0
while row < qweight.shape[0]:
if self.bits in [2, 4, 8]:
for j in range(i, i + (32 // self.bits)):
qweight[row] |= intweight[j] << (self.bits * (j - i))
i += 32 // self.bits
row += 1
else:
raise NotImplementedError("Only 2,4,8 bits are supported.")
qweight = qweight.astype(np.int32)
self.qweight = torch.from_numpy(qweight)
zeros -= 1
zeros = zeros.numpy().astype(np.uint32)
qzeros = np.zeros(
(zeros.shape[0], zeros.shape[1] // 32 * self.bits), dtype=np.uint32
)
i = 0
col = 0
while col < qzeros.shape[1]:
if self.bits in [2, 4, 8]:
for j in range(i, i + (32 // self.bits)):
qzeros[:, col] |= zeros[:, j] << (self.bits * (j - i))
i += 32 // self.bits
col += 1
else:
raise NotImplementedError("Only 2,4,8 bits are supported.")
qzeros = qzeros.astype(np.int32)
self.qzeros = torch.from_numpy(qzeros)
def forward(self, x):
out_shape = x.shape[:-1] + (self.outfeatures,)
out = QuantLinearFunction.apply(
x.reshape(-1, x.shape[-1]),
self.qweight,
self.scales,
self.qzeros,
self.g_idx,
self.bits,
self.maxq,
)
out = out + self.bias if self.bias is not None else out
return out.reshape(out_shape)

83
server/text_generation_server/utils/logits_process.py
View File

@ -1,6 +1,5 @@
import math import math
import torch import torch
import habana_frameworks.torch.core as htcore
from loguru import logger from loguru import logger
from typing import Dict, Union from typing import Dict, Union
@ -44,32 +43,38 @@ class StaticWarper:
if typical_p is not None and typical_p < 1.0: if typical_p is not None and typical_p < 1.0:
self.warpers.append(TypicalLogitsWarper(mass=typical_p)) self.warpers.append(TypicalLogitsWarper(mass=typical_p))
self.hpu_graph = None self.cuda_graph = None
self.static_scores = None self.static_scores = None
self.static_warped_scores = None self.static_warped_scores = None
self.static_next_logprob = None self.static_next_logprob = None
def __call__(self, scores): def __call__(self, scores):
if self.hpu_graph is None: if torch.cuda.is_available():
self.static_scores = scores.clone().contiguous() if self.cuda_graph is None:
self.static_warped_scores = scores.clone().contiguous() self.static_scores = scores
self.static_next_logprob = scores.clone().contiguous() self.cuda_graph = torch.cuda.CUDAGraph()
self.hpu_graph = htcore.hpu.HPUGraph()
with htcore.hpu.graph(self.hpu_graph): with torch.cuda.graph(self.cuda_graph, pool=mempool):
local_scores = self.static_scores local_scores = self.static_scores
for warper in self.warpers: for warper in self.warpers:
local_scores = warper(None, local_scores) local_scores = warper(None, local_scores)
self.static_warped_scores.copy_(local_scores) self.static_warped_scores = local_scores
# Compute logprobs # Compute logprobs
self.static_next_logprob.copy_(torch.log_softmax(self.static_warped_scores, -1)) self.static_next_logprob = torch.log_softmax(
self.static_warped_scores, -1
)
self.static_scores.copy_(scores) self.static_scores.copy_(scores)
self.hpu_graph.replay() self.cuda_graph.replay()
return self.static_warped_scores, self.static_next_logprob return self.static_warped_scores, self.static_next_logprob
# CPU branch
for warper in self.warpers:
scores = warper(None, scores)
return scores, torch.log_softmax(scores, -1)
@lru_cache(10) @lru_cache(10)
def static_warper( def static_warper(
@ -78,7 +83,9 @@ def static_warper(
top_p: Optional[float], top_p: Optional[float],
typical_p: Optional[float], typical_p: Optional[float],
) -> StaticWarper: ) -> StaticWarper:
return StaticWarper(temperature=temperature, top_k=top_k, top_p=top_p, typical_p=typical_p) return StaticWarper(
temperature=temperature, top_k=top_k, top_p=top_p, typical_p=typical_p
)
class HeterogeneousRepetitionPenaltyLogitsProcessor(LogitsProcessor): class HeterogeneousRepetitionPenaltyLogitsProcessor(LogitsProcessor):
@ -95,13 +102,17 @@ class HeterogeneousRepetitionPenaltyLogitsProcessor(LogitsProcessor):
def __init__(self, penalty: List[float], dtype: torch.dtype, device: torch.device): def __init__(self, penalty: List[float], dtype: torch.dtype, device: torch.device):
self.penalty = penalty self.penalty = penalty
self.penalty_tensor = torch.tensor(penalty, dtype=dtype, device=device).unsqueeze(1) self.penalty_tensor = torch.tensor(
penalty, dtype=dtype, device=device
).unsqueeze(1)
def __call__(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.Tensor: def __call__(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.Tensor:
score = torch.gather(scores, 1, input_ids) score = torch.gather(scores, 1, input_ids)
# if score < 0 then repetition penalty has to be multiplied to reduce the previous token probability # if score < 0 then repetition penalty has to be multiplied to reduce the previous token probability
score = torch.where(score < 0, score * self.penalty_tensor, score / self.penalty_tensor) score = torch.where(
score < 0, score * self.penalty_tensor, score / self.penalty_tensor
)
scores.scatter_(1, input_ids, score) scores.scatter_(1, input_ids, score)
return scores return scores
@ -159,11 +170,9 @@ class HeterogeneousFrequencyPenaltyLogitsProcessor(LogitsProcessor):
vocab_size = scores.size(1) vocab_size = scores.size(1)
# Calculate the frequency for each token so far # Calculate the frequency for each token so far
token_freq = torch.zeros( token_freq = torch.zeros(batch_size, vocab_size, device=input_ids.device)
batch_size, vocab_size, dtype=scores.dtype, device=scores.device
)
token_freq.scatter_add_( token_freq.scatter_add_(
1, input_ids, torch.ones_like(input_ids, dtype=scores.dtype, device=scores.device) 1, input_ids, torch.ones_like(input_ids, dtype=torch.float)
) )
token_freq /= input_size token_freq /= input_size
@ -190,9 +199,13 @@ class HeterogeneousTemperatureLogitsWarper:
The value used to module the logits distribution. The value used to module the logits distribution.
""" """
def __init__(self, temperature: List[float], dtype: torch.dtype, device: torch.device): def __init__(
self, temperature: List[float], dtype: torch.dtype, device: torch.device
):
self.temperature = temperature self.temperature = temperature
self.temperature_tensor = torch.tensor(temperature, dtype=dtype, device=device).unsqueeze(1) self.temperature_tensor = torch.tensor(
temperature, dtype=dtype, device=device
).unsqueeze(1)
def __call__(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.Tensor: def __call__(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.Tensor:
scores.div_(self.temperature_tensor) scores.div_(self.temperature_tensor)
@ -231,7 +244,9 @@ class HeterogeneousTopPLogitsWarper(LogitsWarper):
min_tokens_to_keep: int = 1, min_tokens_to_keep: int = 1,
): ):
self.top_p = top_p self.top_p = top_p
self.top_p_opposite = 1 - torch.tensor(top_p, dtype=dtype, device=device).unsqueeze(1) self.top_p_opposite = 1 - torch.tensor(
top_p, dtype=dtype, device=device
).unsqueeze(1)
self.filter_value = filter_value self.filter_value = filter_value
self.min_tokens_to_keep = min_tokens_to_keep self.min_tokens_to_keep = min_tokens_to_keep
@ -248,7 +263,9 @@ class HeterogeneousTopPLogitsWarper(LogitsWarper):
sorted_indices_to_remove[..., -self.min_tokens_to_keep :] = 0 sorted_indices_to_remove[..., -self.min_tokens_to_keep :] = 0
# scatter sorted tensors to original indexing # scatter sorted tensors to original indexing
indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) indices_to_remove = sorted_indices_to_remove.scatter(
1, sorted_indices, sorted_indices_to_remove
)
warped_scores = scores.masked_fill_(indices_to_remove, self.filter_value) warped_scores = scores.masked_fill_(indices_to_remove, self.filter_value)
return warped_scores return warped_scores
@ -296,7 +313,9 @@ class HeterogeneousTopKLogitsWarper(LogitsWarper):
disabled = [x == 0 for x in top_k] disabled = [x == 0 for x in top_k]
if any(disabled): if any(disabled):
self.top_k_disabled_mask = torch.tensor(disabled, dtype=torch.bool, device=device).view(-1, 1) self.top_k_disabled_mask = torch.tensor(
disabled, dtype=torch.bool, device=device
).view(-1, 1)
else: else:
self.top_k_disabled_mask = None self.top_k_disabled_mask = None
@ -332,7 +351,9 @@ class HeterogeneousTopKLogitsWarper(LogitsWarper):
self.max_top_k = max(self.top_k) self.max_top_k = max(self.top_k)
if self.top_k_disabled_mask is not None: if self.top_k_disabled_mask is not None:
self.top_k_disabled_mask = self.top_k_disabled_mask[indices] if any(disabled) else None self.top_k_disabled_mask = (
self.top_k_disabled_mask[indices] if any(disabled) else None
)
return self return self
return None return None
@ -398,11 +419,15 @@ class HeterogeneousTypicalLogitsWarper(LogitsWarper):
if self.disabled_mask is not None: if self.disabled_mask is not None:
last_ind.masked_fill_(self.disabled_mask, scores.shape[-1] - 1) last_ind.masked_fill_(self.disabled_mask, scores.shape[-1] - 1)
sorted_indices_to_remove = sorted_scores > sorted_scores.gather(1, last_ind.view(-1, 1)) sorted_indices_to_remove = sorted_scores > sorted_scores.gather(
1, last_ind.view(-1, 1)
)
if self.min_tokens_to_keep > 1: if self.min_tokens_to_keep > 1:
# Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below) # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
sorted_indices_to_remove[..., : self.min_tokens_to_keep] = 0 sorted_indices_to_remove[..., : self.min_tokens_to_keep] = 0
indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) indices_to_remove = sorted_indices_to_remove.scatter(
1, sorted_indices, sorted_indices_to_remove
)
warped_scores = scores.masked_fill_(indices_to_remove, self.filter_value) warped_scores = scores.masked_fill_(indices_to_remove, self.filter_value)
@ -416,7 +441,9 @@ class HeterogeneousTypicalLogitsWarper(LogitsWarper):
self.mass_tensor = self.mass_tensor[indices] self.mass_tensor = self.mass_tensor[indices]
if self.disabled_mask is not None: if self.disabled_mask is not None:
self.disabled_mask = self.disabled_mask[indices] if any(disabled) else None self.disabled_mask = (
self.disabled_mask[indices] if any(disabled) else None
)
return self return self
return None return None
@ -553,7 +580,7 @@ class HeterogeneousGrammarLogitProcessor(LogitsProcessor):
mask = torch.full_like(logits, -math.inf) mask = torch.full_like(logits, -math.inf)
for i in range(logits.shape[0]): for i in range(logits.shape[0]):
fsm = self.fsms[i] fsm = self.fsms[i]
if fsm is None: if fsm_grammar_states[i] == -1 or fsm is None:
continue continue
allowed_tokens = fsm.allowed_token_ids(fsm_grammar_states[i]) allowed_tokens = fsm.allowed_token_ids(fsm_grammar_states[i])
mask[i, allowed_tokens] = 0 mask[i, allowed_tokens] = 0

125
server/text_generation_server/utils/tokens.py
View File

@ -1,5 +1,3 @@
# Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.
import re import re
from typing import List, Optional, Tuple, Set, Union from typing import List, Optional, Tuple, Set, Union
@ -22,22 +20,21 @@ from text_generation_server.utils.logits_process import (
) )
from text_generation_server.utils.watermark import WatermarkLogitsProcessor from text_generation_server.utils.watermark import WatermarkLogitsProcessor
from transformers import PreTrainedTokenizerBase, RepetitionPenaltyLogitsProcessor from transformers import PreTrainedTokenizerBase, RepetitionPenaltyLogitsProcessor
from transformers.utils import to_py_obj
class NextTokenChooser: class NextTokenChooser:
def __init__( def __init__(
self, self,
watermark=False, watermark: bool = False,
temperature=1.0, temperature: float = 1.0,
repetition_penalty=1.0, repetition_penalty: float = 1.0,
frequency_penalty=0.0, frequency_penalty: float = 0.0,
top_k=None, top_k: Optional[int] = None,
top_p=None, top_p: Optional[float] = None,
typical_p=None, typical_p: Optional[float] = None,
do_sample=False, do_sample: bool = False,
seed=0, seed: int = 0,
device="cpu", device: str = "cpu",
tokenizer: Optional[PreTrainedTokenizerBase] = None, tokenizer: Optional[PreTrainedTokenizerBase] = None,
grammar: str = "", grammar: str = "",
grammar_type: GrammarType = GrammarType.GRAMMAR_TYPE_NONE, grammar_type: GrammarType = GrammarType.GRAMMAR_TYPE_NONE,
@ -70,7 +67,9 @@ class NextTokenChooser:
or (typical_p is not None and typical_p < 1.0) or (typical_p is not None and typical_p < 1.0)
) )
if has_warpers: if has_warpers:
self.static_warper = static_warper(temperature=temperature, top_k=top_k, top_p=top_p, typical_p=typical_p) self.static_warper = static_warper(
temperature=temperature, top_k=top_k, top_p=top_p, typical_p=typical_p
)
else: else:
self.static_warper = None self.static_warper = None
@ -249,12 +248,10 @@ class HeterogeneousNextTokenChooser:
tokenizer: PreTrainedTokenizerBase, tokenizer: PreTrainedTokenizerBase,
grammars: List[str], grammars: List[str],
grammar_types: List[int], grammar_types: List[int],
fsm_grammar_states: List[int], fsm_grammar_states=List[int],
quantization_enabled: bool,
): ):
warpers = [] warpers = []
# TODO: enable watermark with FP8 quantization
self.watermark_processor = ( self.watermark_processor = (
HeterogeneousProcessorWrapper( HeterogeneousProcessorWrapper(
{ {
@ -263,7 +260,7 @@ class HeterogeneousNextTokenChooser:
if do_watermark if do_watermark
} }
) )
if any(watermark) and not quantization_enabled if any(watermark)
else None else None
) )
@ -435,18 +432,6 @@ class HeterogeneousNextTokenChooser:
) )
return self return self
def advance_grammar_single_with_past_state(
self, grammar_state_index: int, next_id: torch.Tensor, past_state: int
):
if self.grammar_processor is not None:
next_id = next_id.item()
self.fsm_grammar_states[grammar_state_index] = (
self.grammar_processor.advance_at_index(
next_id, past_state, grammar_state_index,
)
)
return self
def filter(self, indices): def filter(self, indices):
if self.watermark_processor is not None: if self.watermark_processor is not None:
self.watermark_processor = self.watermark_processor.filter(indices) self.watermark_processor = self.watermark_processor.filter(indices)
@ -497,7 +482,6 @@ class HeterogeneousNextTokenChooser:
device: torch.device, device: torch.device,
tokenizer: PreTrainedTokenizerBase, tokenizer: PreTrainedTokenizerBase,
fsm_grammar_states: Optional[List[int]] = None, fsm_grammar_states: Optional[List[int]] = None,
quantization_enabled: bool = False,
) -> "HeterogeneousNextTokenChooser": ) -> "HeterogeneousNextTokenChooser":
return HeterogeneousNextTokenChooser( return HeterogeneousNextTokenChooser(
watermark=[pb_.watermark for pb_ in pb], watermark=[pb_.watermark for pb_ in pb],
@ -517,37 +501,12 @@ class HeterogeneousNextTokenChooser:
fsm_grammar_states=( fsm_grammar_states=(
fsm_grammar_states if fsm_grammar_states else [0] * len(pb) fsm_grammar_states if fsm_grammar_states else [0] * len(pb)
), ),
quantization_enabled=quantization_enabled,
) )
def pad_next_token_chooser_parameters(
parameters: List[generate_pb2.NextTokenChooserParameters],
expected_size: int,
) -> List[generate_pb2.NextTokenChooserParameters]:
# disable all logits processors to minimize padding overhead
empty_parameters = generate_pb2.NextTokenChooserParameters(
temperature=1.0,
top_k=0,
top_p=1.0,
typical_p=1.0,
do_sample=False,
seed=0,
repetition_penalty=1.0,
frequency_penalty=0.0,
watermark=False,
grammar="",
grammar_type=0,
)
parameters.extend(
[empty_parameters] * (expected_size - len(parameters))
)
return parameters
class Sampling: class Sampling:
def __init__(self, seed: int, device: str = "cpu"): def __init__(self, seed: int, device: str = "cpu"):
self.generator = torch.Generator("cpu") self.generator = torch.Generator(device)
self.generator.manual_seed(seed) self.generator.manual_seed(seed)
self.seed = seed self.seed = seed
@ -583,7 +542,7 @@ class HeterogeneousSampling:
self.greedy = Greedy() self.greedy = Greedy()
def __call__(self, logits): def __call__(self, logits):
out = torch.zeros(logits.shape[0], dtype=torch.int64, device=logits.device) out = torch.empty(logits.shape[0], dtype=torch.int64, device=logits.device)
if self.greedy_indices: if self.greedy_indices:
# Computing for all indices is faster than slicing # Computing for all indices is faster than slicing
torch.argmax(logits, -1, out=out) torch.argmax(logits, -1, out=out)
@ -645,8 +604,9 @@ def batch_top_tokens(
top_n_indices = (logprobs >= nth_highest).nonzero() top_n_indices = (logprobs >= nth_highest).nonzero()
_, top_n_ishes = torch.unique_consecutive(top_n_indices[:, 0], return_counts=True) _, top_n_ishes = torch.unique_consecutive(top_n_indices[:, 0], return_counts=True)
k = 1 if top_n_ishes.numel() == 0 else top_n_ishes.max()
# Take a new topk for these new max n values # Take a new topk for these new max n values
top_k = torch.topk(logprobs, k=top_n_ishes.max(), dim=1, sorted=True) top_k = torch.topk(logprobs, k=k, dim=1, sorted=True)
top_n_ishes = top_n_ishes.tolist() top_n_ishes = top_n_ishes.tolist()
top_indices = top_k.indices.tolist() top_indices = top_k.indices.tolist()
@ -684,50 +644,3 @@ def batch_top_tokens(
batch_top_token_logprobs.append(row_top_token_logprobs) batch_top_token_logprobs.append(row_top_token_logprobs)
return batch_top_token_ids, batch_top_token_logprobs return batch_top_token_ids, batch_top_token_logprobs
def make_tokenizer_optional(tokenizer):
class _(type(tokenizer)):
def __call__(
self,
text,
return_tensors,
padding,
return_token_type_ids,
truncation,
max_length
):
assert return_tensors == "pt", "inccorrect input arguments when calling TransparentTokenizer"
assert padding == "max_length" or padding == "longest", "inccorrect input arguments when calling TransparentTokenizer"
assert return_token_type_ids == False, "inccorrect input arguments when calling TransparentTokenizer"
assert truncation == True, "inccorrect input arguments when calling TransparentTokenizer"
def str_token_to_int(i):
if i == '?':
return tokenizer.pad_token_id
else:
return int(i)
all_tokens = [[str_token_to_int(i.strip()) for i in inner_text.split(',')]
for inner_text in text]
if padding == "longest":
max_length = max(len(tokens) for tokens in all_tokens)
return {"input_ids": torch.tensor([[tokenizer.pad_token_id] * (max_length - len(tokens)) + tokens for tokens in all_tokens]),
"attention_mask": torch.tensor([[0] * (max_length - len(tokens)) + [1] * len(tokens) for tokens in all_tokens])}
def decode(
self,
token_ids,
skip_special_tokens: bool = False,
clean_up_tokenization_spaces: bool = None,
**kwargs,
) -> str:
return ','.join(str(i) for i in to_py_obj(token_ids))
import os
if os.getenv("SKIP_TOKENIZER_IN_TGI", "false").lower() == "true":
tokenizer.__class__ = _
tokenizer.is_transparent = True
def is_tokenizer_transparent(tokenizer):
return hasattr(tokenizer, "is_transparent") and tokenizer.is_transparent is True

12
server/text_generation_server/utils/version.py
View File

@ -1,12 +0,0 @@
from optimum.habana.utils import get_driver_version
from packaging.version import Version
MIN_TGI_GAUDI_SYNAPSE_VERSION=Version("1.16.0")
def is_driver_compatible():
driver_version = get_driver_version()
if driver_version is not None:
if driver_version < MIN_TGI_GAUDI_SYNAPSE_VERSION:
return False
return True

26
server/text_generation_server/utils/watermark.py
View File

@ -34,17 +34,21 @@ class WatermarkLogitsProcessor(LogitsProcessor):
# watermarking parameters # watermarking parameters
self.gamma = gamma self.gamma = gamma
self.delta = delta self.delta = delta
self.rng = torch.Generator(device="cpu") self.rng = torch.Generator(device=device)
self.hash_key = hash_key self.hash_key = hash_key
def _seed_rng(self, input_ids: Union[List[int], torch.LongTensor]): def _seed_rng(self, input_ids: Union[List[int], torch.LongTensor]):
if isinstance(input_ids, list): if isinstance(input_ids, list):
assert len(input_ids) >= 1, "requires at least a 1 token prefix sequence to seed rng" assert (
len(input_ids) >= 1
), "requires at least a 1 token prefix sequence to seed rng"
prev_token = input_ids[-1] prev_token = input_ids[-1]
else: else:
assert len(input_ids) == 1 assert len(input_ids) == 1
input_ids = input_ids[0] input_ids = input_ids[0]
assert input_ids.shape[-1] >= 1, "requires at least a 1 token prefix sequence to seed rng" assert (
input_ids.shape[-1] >= 1
), "requires at least a 1 token prefix sequence to seed rng"
prev_token = input_ids[-1].item() prev_token = input_ids[-1].item()
self.rng.manual_seed(self.hash_key * prev_token) self.rng.manual_seed(self.hash_key * prev_token)
@ -63,7 +67,9 @@ class WatermarkLogitsProcessor(LogitsProcessor):
return greenlist_ids return greenlist_ids
@staticmethod @staticmethod
def _calc_greenlist_mask(scores: torch.FloatTensor, greenlist_token_ids) -> torch.BoolTensor: def _calc_greenlist_mask(
scores: torch.FloatTensor, greenlist_token_ids
) -> torch.BoolTensor:
green_tokens_mask = torch.zeros_like(scores) green_tokens_mask = torch.zeros_like(scores)
green_tokens_mask[-1, greenlist_token_ids] = 1 green_tokens_mask[-1, greenlist_token_ids] = 1
final_mask = green_tokens_mask.bool() final_mask = green_tokens_mask.bool()
@ -76,9 +82,15 @@ class WatermarkLogitsProcessor(LogitsProcessor):
scores[greenlist_mask] = scores[greenlist_mask] + greenlist_bias scores[greenlist_mask] = scores[greenlist_mask] + greenlist_bias
return scores return scores
def __call__(self, input_ids: Union[List[int], torch.LongTensor], scores: torch.FloatTensor) -> torch.FloatTensor: def __call__(
greenlist_ids = self._get_greenlist_ids(input_ids, scores.shape[-1], scores.device) self, input_ids: Union[List[int], torch.LongTensor], scores: torch.FloatTensor
green_tokens_mask = self._calc_greenlist_mask(scores=scores, greenlist_token_ids=greenlist_ids) ) -> torch.FloatTensor:
greenlist_ids = self._get_greenlist_ids(
input_ids, scores.shape[-1], scores.device
)
green_tokens_mask = self._calc_greenlist_mask(
scores=scores, greenlist_token_ids=greenlist_ids
)
scores = self._bias_greenlist_logits( scores = self._bias_greenlist_logits(
scores=scores, greenlist_mask=green_tokens_mask, greenlist_bias=self.delta scores=scores, greenlist_mask=green_tokens_mask, greenlist_bias=self.delta