The LLM Leap: moving a streaming pipeline from small encoders to Gemma 2

Our real-time review enrichment pipelines implemented as an Apache Beam Dataflow job on Google Cloud serves as the cornerstone for many Trustpilot products and services. While effective, we eventually faced a model performance ceiling: our in-house encoder-based model plateaued and proved hard to improve. To overcome this we transitioned to a Gemma 2 9B decoder model maintaining sub-second latency while handling a 30x increase in model size.

Problem: the model performance ceiling

Review enrichment models are part of the streaming job that reads reviews from a PubSub topic and emits sentences enriched with sentiment and topic data into the output queue for downstream consumption and storage.

The job itself has multiple steps and uses in-process models. The sentences are converted to vector embeddings using an in-house transformer based model that runs on Nvidia Tesla T4 GPU followed by sentiment and topic computation. There are also other enrichments happening in the same pipeline that we don’t have alternatives yet, thus making decommissioning of the job impossible just yet.

The bottleneck in the above system is the embedding model that is responsible for “language understanding”. Squeezing additional performance from this legacy model was a dead end and architectural breakthrough was needed while preserving overall system performance and backward compatibility.

Solution: a decoupled Gemma 2 architecture

The breakthrough was provided by the Gemma 2 9B Open Model. Our Data Science team performed fine-tuning of this model for the sentiment analysis and topic modelling tasks, while we (ML Ops) had to productionalise it and integrate into an existing streaming system. This was a challenging task as the model was ~30 times larger than the legacy encoder model.

Our final architecture is shown below. It features:

• the old Dataflow job that produces sentences
• two new lightweight Dataflow jobs to perform enrichments
• set of Vertex Endpoints called Classifier and LLM

In the above diagram “LLM Vertex Endpoint” refers to a vLLM inference engine deployed to Google’s Vertex AI Endpoint. The endpoint serves the fine-tuned model. The “Classifier Vertex Endpoint” refers to a pre/post processing layer. It handles prompt assembly, request conversion, retries with fallback policies (variation of prompt, temperature, etc) and structured output parsing. Due to the fact that fine-tuning was done for two specific tasks we have two pairs of endpoints. The reason we chose Vertex Endpoint was due to convenient integration of endpoints with Apache Beam SDK via VertexAIModelHandlerJSON that handles retries and throttling out of the box.

We arrived at this much larger system incrementally:

1. Attempt 1 (in-place): First we attempted to update the existing Dataflow job in-place and make two chained remote calls to classifiers. However, due to the additional enrichments I mentioned above the embedding model had to stay in place. In our experience, Dataflow jobs that use GPUs take much longer to set up and start processing data (in our case ~40 minutes). This results in a considerable backlog of messages that proved to be impossible to clear out with Vertex endpoints.
2. Attempt 2 (single new enrichment job): Our next step was to refactor Vertex based enrichment into its own Dataflow job. The new job became lightweight and could start processing data in just 5 minutes. Yet due to different computational complexity of sentiment and topics models, autoscaling the job and endpoints in lockstep became challenging eventually causing slow increase in backlog during any spike in the input stream.
3. Attempt 3 (decoupling → success): Thus, we arrived at our final solution which completely decoupled each enrichment type into its own Dataflow job, classifier and LLM endpoints with their specific autoscaling setup.

Challenges

Before facing the above architectural problems we also encountered a number of other issues that slowed down our delivery.

Networking with Private Vertex Endpoints

As per Google’s recommendation we implemented Classifier and LLM as Private Vertex Endpoints with Private Service Connect. However, we discovered an undocumented limitation that does not allow ingress traffic to Private Vertex Endpoint from any other endpoint. After extensive investigation and direct call with our Google support team, we confirmed the limitation and switched to the next recommended setup, Public Dedicated Endpoints.

Vertex AI SDK

Google’s Vertex AI SDK, recommended for interacting with endpoints, has major performance bottlenecks. It relies on synchronous communication which is unacceptable for our IO bound workload. This became a massive problem in our Classifier layer and we addressed it by re-implementing the communication layer using Python’s asyncio library.

GPU availability

With LLMs being used everywhere the GPU availability is a real problem, especially if you are limited to specific deployment regions due to regulations. Our end solution involved using Nvidia A100 40Gb GPUs for LLM Endpoints that had to be deployed in the EU. There is only a single region, europe-west4, with those GPUs and it is really hard to procure them on demand. Our solution was a pragmatic “engineer’s hack”: a Terraform script running in a loop until the compute reservation was successfully provisioned.

Performance

We measured success in Requests Per Second (RPS) rather than tokens per second treating the LLMs as traditional APIs because our sentence statistics were fairly homogeneous. Our target was to achieve 100 RPS with p95 latency under 1 second on a single node LLM endpoint.

For load testing our endpoints we relied on Grafana’s k6 tool. We first tested the LLM to find optimal GPU config (for example, we started off with L4 GPUs) and vLLM settings. Once we achieved target performance for an LLM endpoint we tested the pair Classifier + LLM. This is where we faced the biggest challenge. No matter how many Classifier replicas we created we could not saturate LLM endpoint, i.e. it looked like there is a networking bottleneck in Classifier implementation or Vertex networking. As you can see below, beyond 20 RPS load and throughput diverge and latency sky rockets.

Performance before

Profiling the Fast API application pinpointed the problem, Vertex AI SDK that uses synchronous calls under the hood relying on requests library. So we re-implemented the http communication layer using asyncio library and immediately saw the jump in the performance. With single node LLM Endpoint (a2-highgpu-1g, NVIDIA_TESLA_A100 GPU) and single node classifier (e2-highcpu-4 with 4 Uvicorn workers) we could achieve the below performance of ~300 r/s with p95 latency <1s.

Performance after

Results

Despite the architectural hurdles the migration was a success. Gemma 2 tailored to our tasks achieved Gemini-like precision and recall at a fraction of the cost. Perceived error reduction by F-score for sentiment analysis was 55% and for topic modeling 57%. It allowed infrastructure sovereignty with a smaller environmental footprint while also making future upgrades fairly straightforward.

The project fundamentally leveled up our LLM Ops maturity. Thanks to this project we had an opportunity to collaborate with the Google team. The 4 week engagement process with Google engineers allowed us to vet our architecture, improve load testing and CI/CD pipelines. Our experience offered Google direct insight into the real-world performance and limits of their Vertex products. Finally, our generic Classifier layer has unlocked a roadmap of new LLM use cases, allowing us to onboard both discriminative and generative applications with minimal friction.

Acknowledgements

The project was a massive team effort. We would like to thank Dario Banfi and Michael Cohen Hjertén from Google who helped us to quickly address complex Vertex AI issues alongside the internal Vertex AI team. Special thanks to Subu Ramasubramanian from Trustpilot who organised the Google engagement experience and provided massive support along the way.