Trace

End-to-end Generative Optimization for AI Agents

[This repository accomponanies the Trace paper. It is a fully functional implementation of the platform for generative optimization described in the paper, and contains code necessary to reproduce the experiments reported. This library was implemented and maintained by the authors while they were at Microsoft.]

Trace is a new AutoDiff-like tool for training AI systems end-to-end with general feedback (like numerical rewards or losses, natural language text, compiler errors, etc.). Trace generalizes the back-propagation algorithm by capturing and propagating an AI system's execution trace. Trace is implemented as a PyTorch-like Python library. Users write Python code directly and can use Trace primitives to optimize certain parts, just like training neural networks!

Paper | Project website | Documentation | Blogpost | Discord channel | Mailing list

Setup

Simply run

pip install trace-opt

Or for development, clone the repo and run the following.

pip install -e .

The library requires Python >= 3.9. By default (starting with v0.1.3.5), we use LiteLLM as the backend of LLMs. For backward compatibility, we provide backend-support with AutoGen; when installing, users can add [autogen] tag to install a compatible AutoGen version (e.g., pip install trace-opt[autogen]). You may require Git Large File Storage if git is unable to clone the repository.

For questions or reporting bugs, please use Github Issues or post on our Discord channel. We actively check these channels.

Updates

• 2025.2.7 Trace was featured in the G-Research NeurIPS highlight by the Science Director Hugh Salimbeni.
• 2024.12.10 Trace was demoed in person at NeurIPS 2024 Expo.
• 2024.11.05 Ching-An Cheng gave a talk at UW Robotics Colloquium on Trace: video.
• 2024.10.21 New paper by Nvidia, Stanford, Visa, & Intel applies Trace to optimize for mapper code of parallel programming (for scientific computing and matrix multiplication). Trace (OptoPrime) learns code achieving 1.3X speed up under 10 minutes, compared to the code optimized by a system engineer expert.
• 2024.9.30 Ching-An Cheng gave a talk to the AutoGen community: link.
• 2024.9.25 Trace Paper is accepted to NeurIPS 2024.
• 2024.9.14 TextGrad is available as an optimizer in Trace.
• 2024.8.18 Allen Nie gave a talk to Pasteur Labs & Institute for Simulation Intelligence.

QuickStart

Trace has two primitives: node and bundle. node is a primitive to define a node in the computation graph. bundle is a primitive to define a function that can be optimized.

from opto.trace import node

x = node(1, trainable=True)
y = node(3)
z = x / y
z2 = x / 3  # the int 3 would be converted to a node automatically

list_of_nodes = [x, node(2), node(3)]
node_of_list = node([1, 2, 3])

node_of_list.append(3)

# easy built-in computation graph visualization
z.backward("maximize z", visualize=True, print_limit=25)

Once a node is declared, all the following operations on the node object will be automatically traced. We built many magic functions to make a node object act like a normal Python object. By marking trainable=True, we tell our optimizer that this node's content can be changed by the optimizer.

For functions, Trace uses decorators like @bundle to wrap over Python functions. A bundled function behaves like any other Python function.

from opto.trace import node, bundle

@bundle(trainable=True)
def strange_sort_list(lst):
    '''
    Given list of integers, return list in strange order.
    Strange sorting, is when you start with the minimum value,
    then maximum of the remaining integers, then minimum and so on.
    '''
    lst = sorted(lst)
    return lst

test_input = [1, 2, 3, 4]
test_output = strange_sort_list(test_input)
print(test_output)

Now, after declaring what is trainable and what isn't, and use node and bundle to define the computation graph, we can use the optimizer to optimize the computation graph.

…

Then, we can use the familiar PyTorch-like syntax to conduct the optimization.

Here is another example of a simple sales agent:

…

Imagine we have a feedback function (like a reward function) that tells us how well the agent is doing. We can then optimize this agent online:

…

Defining and training an agent through Trace will give you more flexibility and control over what the agent learns.

Tutorials

Supported Optimizers

Currently, we support three optimizers:

• OPRO: Large Language Models as Optimizers
• TextGrad: TextGrad: Automatic "Differentiation" via Text
• OptoPrime: Our proposed algorithm -- using the entire computational graph to perform parameter update. It is 2-3x faster than TextGrad.

Using our framework, you can seamlessly switch between different optimizers:

optimizer1 = OptoPrime(strange_sort_list.parameters())
optimizer2 = OPRO(strange_sort_list.parameters())
optimizer3 = TextGrad(strange_sort_list.parameters())

Here is a summary of the optimizers:

The table evaluates the frameworks in the following aspects:

• Computation Graph: Whether the optimizer leverages the computation graph of the workflow.
• Code as Functions: Whether the framework allows users to write actual executable Python functions and not require users to wrap them in strings.
• Library Support: Whether the framework has a library to support the optimizer.
• Speed: TextGrad is about 2-3x slower than OptoPrime (Trace). OPRO has no concept of computational graph, therefore is very fast.
• Large Graph: OptoPrime (Trace) represents the entire computation graph in context, therefore, might have issue with graphs that have more than hundreds of operations. TextGrad does not have the context-length issue, however, might be very slow on large graphs.

We provide a comparison to validate our implementation of TextGrad in Trace:

To produce this table, we ran the TextGrad pip-installed repo on 2024-10-30, and we also include the numbers reported in the TextGrad paper. The LLM APIs are called around the same time to ensure a fair comparison. TextGrad paper's result was reported in 2024-06.

You can also easily implement your own optimizer that works directly with TraceGraph (more tutorials on how to work with TraceGraph coming soon).

LLM API Setup

Currently we rely on LiteLLM or AutoGen v0.2 for LLM caching and API-Key management.

By default, LiteLLM is used. To change the default backend, set the environment variable TRACE_DEFAULT_LLM_BACKEND on terminal

export TRACE_DEFAULT_LLM_BACKEND=""  # 'LiteLLM' or 'AutoGen`

or in python before importing opto

import os
os.environ["TRACE_DEFAULT_LLM_BACKEND"] = ""  # 'LiteLLM' or 'AutoGen`
import opto

Using LiteLLM as Backend

Set the keys as the environment variables, following the documentation of LiteLLM. For example,

import os
os.environ["OPENAI_API_KEY"] = ""
os.environ["ANTHROPIC_API_KEY"] = ""

In Trace, we add another environment variable TRACE_LITELLM_MODEL to set the default model name used by LiteLLM for convenience, e.g.,

export TRACE_LITELLM_MODEL='gpt-4o'

will set all LLM instances in Trace to use gpt-4o by default.

Using AutoGen as Backend

First install Trace with autogen flag, `p