How Flo Health built safer, more useful open models with Databricks AI Runtime

來源原標題

How Flo Health built safer, more useful open models with Databricks AI Runtime

機構

Databricks

原始標籤

Company · Customers

原始連結

https://www.databricks.com/blog/how-flo-health-built-safer-more-useful-open-models-databricks-ai-runtime

摘要原文

For Flo Health, the bar for generative AI is unusually high. With 81 million monthly active users, the women’s health app operates at a scale where even small weaknesses in an AI system matter. Because Flo’s products deal with health and well-being, responses need to be medically safe and genuinely useful. As Flo moved from using only traditional ML towards personalized insights powered by generative AI, the team saw the potential for a more flexible, personalized product. But relying on a single general-purpose model introduced other constraints: cost at scale, dependence on a model provider’s lifecycle, and limited ability to encode Flo’s own domain expertise. Flo’s longer-term vision was a compound AI system that could use the right model for each task while maintaining a consistent safety standard across them. That led Flo toward fine-tuned open models and a new post-training pipeline built with Databricks. The team used AI judges and privacy-first synthetic data to turn Flo’s standards for medical safety and usefulness into a training signal, then used Databricks AI Runtime to run distributed fine-tuning on serverless NVIDIA H100 GPUs. That approach also made sense economically at Flo’s scale. The team estimates that using tailored, fine-tuned models can reduce inference costs by an order of magnitude compared with relying on frontier model APIs, while giving Flo greater control over model behavior and lifecycle. But before Flo could train models against its own standards, it needed a reliable way to define and measure what “good” looked like. Flo defined two goals for its AI systems: medical safety as the foundation and usefulness as the goal. A model can avoid unsafe statements and still provide little value, or give a detailed, personalized answer while missing an important medical safeguard. Consider a question such as, “Why is my period late?” An answer that simply tells the user to see a medical professional may be safe but not particularly useful. An answer that lists only common lifestyle causes may be more helpful, but incomplete from a medical-safety perspective. Flo wanted responses that could balance both: useful information, appropriate context, and the right safety netting. To make that balance measurable, Flo’s medical, product, and engineering teams built an evaluation framework that decomposed medical risk into specific rubrics that could be tested consistently. For example, after identifying that responses to symptom questions could be overly reassuring, Flo created a “balanced causes” rule requiring appropriate coverage of both benign and potentially serious causes when relevant. Those rubrics were then turned into AI judges, each with a definition, grading instructions, and training cases. Flo validated the judges against synthetic cases and anonymized cases from real-user experiments. When an AI judge disagreed with an initial human assessment, blinded clinical reviewers adjudicated the disagreement to establish ground truth. Flo then used part of those cases to refine the judges and reserved the rest to measure performance. With the evaluation system in place, Flo could use those same judges as the reward signal for training. The final framework included 26 judges across three categories: Building a conventional human-labeled fine-tuning dataset would have required significant clinical expertise for every example. Flo concentrated its medical experts’ time on designing and calibrating the evaluation system, which could then score training candidates at scale. The resulting pipeline generated more than 15,000 synthetic questions using privacy-first principles and created training pairs without human-annotated responses. Flo derived the inputs from model-generated or anonymously aggregated signals rather than user-written free text or identifiable information. Clinicians remained essential to defining and validating what “safe” meant, but they no longer needed to manually label every individual training pair. With calibrated judges and synthetic training data in place, Flo adopted a Reinforcement Fine-Tuning (RTF) approach that generates multiple candidate responses during training, uses the evaluation system to identify the strongest ones, and fine-tunes the model so that quality is encoded in its weights rather than relying on expensive best-of-N search on every production request. Flo’s workflow followed five stages: The first version of this approach exposed an important limitation: usefulness improved while medical safety remained relatively flat because the preference pairs did not distinguish sufficiently between degrees of medical quality. Flo addressed that by adding a graded safety score. Responses that passed the binary safety gate were re-ranked according to how thoroughly they met the applicable medical criteria, creating preference pairs with a clearer medical-safety contrast. Running this generate-score-select-train loop at scale is compute-intensive. Flo’s training stack combined Databricks AI Runtime, LLM Foundry, and Compose-RL, with MLflow used for tracking and checkpoints. AI Runtime provided the serverless GPU environment, LLM Foundry handled distributed training, and Compose-RL provided the preference-optimization layer. For its fine-tuning run, Flo used 32 NVIDIA H100 GPUs with PyTorch and Fully Sharded Data Parallel (FSDP), and ran SFT followed by DPO via LLM Foundry and Compose-RL. That setup gave engineers on-demand GPU compute for distributed deep learning training, allowing them to move from synthetic data generation to evaluation to fine-tuning without manually standing up and configuring GPU clusters for each experiment. Keeping the surrounding data, experimentation, and ML tooling in the same environment made it easier to iterate across the full post-training workflow. That workflow produced measurable gains across both dimensions Flo cared about. The first model Flo evaluated scored 34% on its medical-safety benchmark and 42% on usefulness, while only 16% of responses met both criteria simultaneously. After fine-tuning, medical safety increased to approximately 61%, usefulness reached 76%, and the combined score climbed from 16% to 47% — nearly three times the starting point. The 47% combined score was an important milestone, but the more durable outcome was the reusable evaluation and training system Flo had built around it. Flo can carry that system forward as the underlying models evolve. The team plans to apply it to newer open-weight models, including Gemma, Qwen, and Kimi, while preserving the medical safety and usefulness criteria it has already developed and validated. With Databricks AI Runtime underpinning that loop, Flo can iterate on new models without rebuilding the training environment, keeping its engineers focused on medical expertise, model quality, and the user experience. Subscribe to our blog and get the latest posts delivered to your inbox.