From Gradients To Governance: Making PyTorch Lineage-Aware - Kateryna Romashko & Clodagh Walsh

This talk addresses the critical need for lineage awareness in PyTorch as it transitions into enterprise environments that are subject to stringent regulations like the EU AI Act and GDPR. The speaker discusses how current AI projects often overlook tracking the provenance of training data, leading to compliance risks, data leaks, and cultural incompatibilities. They advocate for integrating lineage tracking directly into PyTorch to ensure that data usage complies with legal and ethical standards. Open Lineage is proposed as a solution to capture metadata about data flows and model training events. The speaker highlights existing tools and the limitations they pose, urging the development of a unified, non-invasive observability framework to enhance accountability in machine learning workflows. The session concludes with an exploration of potential implementation strategies for incorporating lineage tracking within the PyTorch ecosystem.

Hi everyone. Thank you for joining for our talk. Today we're going to talk about uh from gradients to governments making PyTorch lineage aware. Um Kateryna Romashka, uh a senior software engineer at Red Hat emerging technology team. And I'm joined today by my colleague Clodagh Walsh, software engineer, who will take over um for the second half of the talk. So, we are here to talk about something that

has been keeping AI engineers awake at night. Not whether their model converges, but whether they allowed to train in the first place. Over the next 25 minutes, we're going to argue that lineage isn't just metadata you log after training. It's a runtime constraint that belongs inside PyTorch itself. let me start with why this matter right now. Take a look uh sorry. Take a look at uh this

timeline. From 2024 to 2028, two things are happening right now at the same time. So, first, PyTorch is moving from research labs into enterprise production, mission-critical systems in health care, finance, and government. Second, regulators are catching up. So, the EU AI Act enforcement begins in this year. By 2027 and 2028, we enter what I call the compliance deadline zone. Auditors will no longer ask, "Does your model

work?" They will ask, "What data trained this model? Where did that data flow? And was that flow allowed?" So, let me break down what this means for PyTorch deployments. So, first of all, the um regulatory in p- imperative. So, GDPR, um LGPD, and the EU AI Act all demand complete auditability. They require constant tracking. They require policy enforcement. There is aren't there is aren't suggestions. They are

legal requirements with the fines attached. So, that is no longer static. Your training data doesn't sit in one store anymore. It flows across cloud providers, geographic regions, and distributed system in real time. By the time you have trained one model, your data may have crossed three borders. 70% prioritize sovereignty. According to the recent enterprise service, 70% of organization now mandate data sovereignty as a principal element in

their cloud architecture decisions. This isn't a niche concern. It's a part on priority. The critical gap Here is a nuance most people miss. Data residence, keeping data in a region is not the same as data sovereignty. Sovereignty means governance flows data wherever it flows. If your model trains on EU data, but runs on US infrastructure, residence law doesn't protect you. So, why should PyTorch care? Because PyTorch

is where the training happens. If lineage isn't tracked inside the runtime, it's not track tracked at all. So, the problem isn't in your code. Your training loop can be perfect. Your hyperparameters optional. Your model state of the art. The problem is and the origin of your data. Let me give you three real world example of lineage failures. First, data leak. Remember the ChatGPT accident where private chats

were accidentally published? The root cause wasn't a bug in the the model architecture. It was in lineage tracking fail. The model the model couldn't distinguish between private conversation context and public context. It didn't know which data came from where. Second, invalid output. We are seeing models trained on web scraped data start to hallucinate fake laws and regulations. Why? Because the training included random forums forums and blog

posts analyze actual legal documents. Without lineage, the model has no proof of which source authoritative. It treats everything equally. So, the third culture incompatibility. A model trained primarily on the West data might produce outputs that violate cultural norms in, for example, Southeast Asia or the Middle East. The problem isn't be biased in the traditional sense. It's that the training data lacks provenance markers for cultural and legal

context. Here's what these failures have in common have in common. They all happen because data flows across borders, but policies don't travel with data. So, you can see this visually on the right. Data originate in one jurisdiction, gets processed in another, and somewhere along the way the governance the governance context the current context gets lost. This isn't theoretical. There is a production incident that cost companies millions

and millions in finances and reputation damage. So, the question becomes how do we fix this at the framework level? So, here's the core of our argument. PyTorch brilliantly tracks how models learn. Look at what it give us today. Gradients, computer graphs, training loss, model parameters. These primitives optimize learning. They make training faster, more stable, more effective. But PyTorch does not track where models should learn from specific

data. what's missing then? Data origin. Which country, which data set, which version did this tensor come from? Consent boundaries. Is it that a law for training or only for inference? Has consent been withdrawn? Regulatory jurisdiction. Does GDPR apply or any other policy? Purpose limitations. Is this data restricted to non-commercial use or only for research? So, let me show you why existing tools can't deliver this. Yeah. the

industry has tried two domain approaches to lineage. Both have the same blind spot. First approach, API based instrumentation. Tools like MLflow, TFX, and Kubeflow. So, from the good side, they are flexible. You log what you want. Uh from other side, they invasive. You have to refactor your training code. You add logging calls everywhere. And because it's manual, it's error-prone. Users forget to log um artifacts. They log

um inconsistently across pipeline stage. And when you want to switch to a different tracking system, you rewrite your code. Second approach, file system metadata extensions. So, capture happens automatically at the file system level. Metadata replica um replication is consistent, but at the same time you're tied to the storage layer. You can capture computational lineage, what happened inside the training loop, and it doesn't span across different systems.

So, that's the that's the bigger problem cuz all these tools they work independently and they all have uh like connection. They So, they have a connection, but they isolated from each other. So, for example, when you try to answer a question like, "Which Spark job prepared data that went into the this Kubeflow component, which triggered this MLflow training run?" You just can't answer. There's no common language.

Uh what we need first is some make a standard, non-invasive, um end-to-end of observability. A way to see the full uh lineage of an ML pipeline from data injection to model without writing every script. So, that's exactly what Open Lineage provides. Open Lineage is an open standard for collecting lineage metadata. It's a framework agnostic link and language agnostic. Where Marquez is a popular reference implementation. What does

Open Lineage give you? So, visibility into jobs, datasets, and their relationships. A standard schema for audit trails and reporting. A foundation for lineage world tooling across your entire pipeline. You can see the architecture here. So, on the bottom we have um ki- kind of producer uh producer tools. Uh training So, they training jobs and Open Lineage events, uh which go to Marquez, and then you can query

answer to question like what data set trained this model. So, this is powerful. And now, so let me show you uh what Open Lineage looks like in practice. Not as theory, but as something we actually integrated. So, on this slide, you see in three columns. They are from our project, where we connected Open Lineage across ML pipeline. First column, MLflow. So, normally MLflow tracks parameters and metrics,

but with the Open Lineage tracking store, every MLflow dot start run automatically emits a start event to Marquez. Every log param and log metrics gets accumulated. And when you're calling end run, it emits a complete event with everything captured. So, no extra code, and the user even doesn't know um there's Open Lineage. So, second column, KubeFlow. Inside each pipeline component, we wrap the the work the work.

Um on enter, it emits start, on exit, complete or fail. Um it parses the input and output data sets your eyes, and turns them into Open Lineage data sets automatically. On the third column, you can see Marquez, which is also uh some UI of Open Lineage. Um So, it's uh as I said recently, Marquez is uh reference implementation that ingests all these events. It builds a graph,

and on the bottom, you can see uh just a screenshot, like little part from our project, how it looks in the in real life. Uh so, that's how you can follow from origin data set, some job was made on the data set, and we have output data set, and have connection between them. So, what do we get from this? Uh we get observability. We can answer which

data sets fed into which Kubeflow component, which MLflow run produced that model, what hyper parameters were used. But, here is a critical distinction, and I want to be very clear about this. Open lineage records what happened after it happened. the events uh I emitted post hoc. The training run, uh the model saved, then the events went to Marquez. Then you can query queries lineage. So, this is

excellent for audit trails. Um it's essential for compliance But, now the question is about PyTorch. So, this example how we integrated in our project, and now Cloud going to share some ideas how we see it in Yeah. So, as Kate mentioned, uh we've identified several approaches of how we could implement la- data lineage tracking within PyTorch. Um so, there's two kind of distinct options we've seen, and

then a third open um kind of track that we'll discuss later. Um so, the first and maybe most easiest and naive option would be to implement PyTorch as a uh open lineage producer, much like uh Spark and Airflow data Kate mentioned earlier. So, to do this, we need to make some changes to the data set or data loader components within the PyTorch stack. Uh the second option,

and perhaps the more advanced and um like friendly to the ecosystem, will be to make an internal lineage engine within PyTorch. To do this it obviously have to create the lineage engine itself and then we um, we assume that we need to make some changes to the tensor structure as well. And then so with option C, we are open to ideas from the community. Maybe we need

some new components or edit some components that we haven't yet considered as it's a quite a big ecosystem in stock. So um, the first approach is to make it a producer. Um, so as you can see here we have the open lineage API and the events will be sent to this API which are then looked at via the lineage a visualizer and an auditor can go to

there to see all the events that happen. So we start at the start of your training pipeline which is when you bring in the data set um, into your Jupyter notebook or whatever your environment is where you're training your data and your this usually starts with the data loader component in PyTorch. Um, you also have a data set then you have batches, samplers and all these kind

of components that help you get your data ready. These can all emit different events at the different levels of granularity that you might need and these um, events will be sent to the open lineage um, API. Uh, so once you have that ready and you're on to the model training step that then produces an artifact. You have the hyper parameters, number of epochs and obviously the data

set that you used. This again will emit another event to the open API open lineage API. Um, after you have the model trained before you save it, you might use some other operations like compression, quantization, etc. Uh, you might want to also track these cuz they might have an effect on your data lineage pipeline. So once you've saved it and done any conversions as well, uh, we

emit another event to the to the system and then they can always be got by the auditor at later on. Um, so with the those events added to PyTorch as like the base level, then we can extend it beyond, um, just emitting events and we can implement like the enforcement and governance layer. So, in the orange above that line there, that's the previous uh, slides and we

have the classic load, train, and save steps. But, the difference here is that instead of emitting the events to the Open Lineage API, we emit them to our internal lineage engine, which, as you can see here, has a policy store that it references. And the policy store is kind of the key to all this because it has, um, the rules, uh, policies, um, any kind of information

that is needed to describe like what the data can do, who can access it, uh, where they can access it from. So, you can think about it like, uh, the or back, if you like. Um, then in the blue stream here, we have the inference, um, inference stage of it. So, a user queries the model or the model or the inference server can query the internal engine.

And then it can see if it's trying to ask a question that references some data that it can't or is trying to, um, ask a question about like an area, um, or a region that it's not supposed to. So, depending on the results of that, they either get a response or some kind of denial message explaining that it's not within their authority to do this. so, that

is the more native model that we, um, see as being more robust and, um, native to PyTorch. Um, so, as I mentioned earlier, we'd also need to modify potentially the tensors within, uh, the PyTorch project to do this. So, when you have a lineage tracking system, you need to track these fields here, like the origin, jurisdiction, the consent terms, and a pointer to the parent. Um, the

pointer to the parent is particularly important because they need to see what happened before and they need to be able to track what's happening afterwards. Um so because time series and PyTorch already have a couple of metadata fields describing the shape, the device, um size, things like that, it seemed a natural fit to put it in there. Um we would suggest that this would be an optional

field because um maybe not every system wants to track lineage, or um there are some space requirements and things like that that you might not want to uh make it a mandatory thing to do. So then, as we were looking into this, uh we came across uh the autograd engine, that's one of the main features of PyTorch. And at first, we kind of thought that this might

be the place uh to implement that lineage. Um with especially that fact that there's the autograd engine makes the gradient graph, and we need a graph for a data lineage. And as tempting as it might be to make the changes in there, uh fundamentally it doesn't make sense because the gradient graph is very different from a data lineage graph. And while you could argue that you could

modify the autograd engine to produce a graph like this, it kind of breaks the fundamental components or fundamental principles, rather, of um good software design, that you shouldn't have too many responsibilities for one particular component. So it makes sense um isolate it into its own thing. And then, another thing to consider is during inference time, the autograd engine is disabled because you don't need to compute the

gradients at that time. So if the engine that has the data lineage isn't used during inference, it's going to be very difficult to try and do the enforcement step. So again, it makes sense um not to use this component. So that's just an aside. Um Then, some other things that we saw along the way was um this thing called uh torch, yes, torch arrow. Um so it's

a preprocessing library that you can use to um kind of edit and modify your data before you throws through the pipeline. Again, because this is where your data is being modified, it makes sense that you might want to track some data lineage events in here. But, um this library is not um uh particularly used right now. It's kind of um a deprecated component. But, we are thinking

that maybe somewhere like this and other components um it could give them like a new lease of life. It gives them a a purpose. So, potentially we could uh look at putting them in here. Then, there's a second component that we came across called Data Pipes. It's a kind of similar idea where you create your own graph of steps of how to create your data set. And

again, it seems like a potential area that we could implement this. But again, it's not um an active part of the system right now, but maybe it could be um looked at again. Uh so, as I mentioned, um it's still kind of like a work in progress. We're thinking about where this fits into the stack. There's many different levels of the stack that needs to go into.

So, we're very open to hearing from the community about your ideas about how you this uh should be done. so, other open questions we have are what would this look like if we had a distributed training environment where you have multiple nodes and clusters. Um also, if you have mixed data sets. So, if you have data coming from multiple different regions, that's a very important use case

to track. Um as I mentioned, there might be potential performance overhead if we're tracking uh the lineage through uh the system. So, if you're, you know, running on small form devices, uh like small language models, it might not be something that you want to do. So, but there could also be things that we could do to um make it more performant. And then, the last thing is

attestation. So, just because we have a data lineage pipeline that shows that this data set originated from some government department, and it looks like it's an official data set. Unless it's signed and verified in some way, we can't really um assume that that is the case. So, that's another thing that we need to consider as part of the Um so, on the screen over here, we have

a QR code to a post that we opened on the Play Torch forums um, where we describe the project a bit and uh, we'd love for you to leave your comments and suggestions there and um, help us build uh, Play Torch and data lineage. Um, yeah, so that's it. Thanks uh, for attending.