What you're asking for is the rough equivalent of asking a C programmer to name their loop variables "index" instead of "i." Everyone familiar with the concepts of c programming knows what "i" means in the context of a for loop. Similarly, everyone familiar with transformers knows what "gelu" and "attn" mean.
They should use proper variable names if they want to have the code understood by anyone non-specialist, and by people who use different terminology, and people looking back at the code in the future when terminology may have changed.
I don’t know about this domain, but the single-letter-variable name etc AKA “match the equation” is a curse when non-CS engineers/scientists write code. It often breaks code conventions, leaving IDEs to light up like a Christmas tree when opening the source. There’s a good reason CS moved from register letters to something closer to natural language.
So then why don't you believe me when I tell you that all of these variable names are extremely standard, and will be familiar to anyone who has written deep learning code before?
I'm a software engineer with a background in ML, so even though I somewhat know the domain language I still get mad at the blatant disrespect for PEP-8. That being said, this one is definitely one of the better codebases I have come across. This feels like it could be fairly easily worked with and understood. I have seen far, far worse code to go along research papers.
These companies hire a lot of engineers straight out of undergrad/master's degrees. The interviews test leetcode knowledge, and today lots of degrees are heavy on Python-scripted ML homework.
The result is companies with billion dollar funding and world-changing goals having a lot of their code look like complete spaghetti.
And this is the engineers who are meant to clean up research scientist code. Scientists generally don't feel like it's their responsibility to write strong code.
Systems-side teams/orgs have better code, but essentially as soon as you enter the 'ml engineer/research engineer/research scientist' layer, it's doomed.
I think this code is actually well-written and maintainable. This is proven in practice because we've adopted it many places in OpenAI, and I've personally found it very easy to adapt to other use-cases (certainly much more so than the from-scratch Transformer implementations I've written!).
As https://news.ycombinator.com/item?id=21456605 points out, the complexity of the code arises from the complexity of the underlying algorithm. Complexity due to software engineering concerns, like Tensorflow scopes, are elegantly handled. [edited for clarity:] Writing a Transformer in 174 lines of code requires a lot of deep thinking about the right underlying abstractions.
> but essentially as soon as you enter the 'ml engineer/research engineer/research scientist' layer, it's doomed.
We actually don't do this! Our only official technical title is "member of technical staff". (People sometimes choose to self-identify as an engineer or researcher, so you might see that on LinkedIn, but we don't have a distinction internally.) Everyone is responsible for their own code, and people care quite a bit about writing code that others can build on.
So I don't have a big problem with some of the function definitions which can be compact, as the other comment points out.
The reason I don't like this code is that it does not comment anything on the critical bits. I don't necessarily care about whether you call the input to your matmul 'x' or 'tensor' or 'input' (although consistency is nice).
The thing that would stop be from absorbing and modifying this code is that it does not comment on all the bits that are non obvious to me if I haven't written a Transformer before. For example:
'Same as tf.matrix_band_part(tf.ones([nd, ns]), -1, ns-nd), but doesn't produce garbage on TPUs.' - I will have to ask the colleague what that means. Why not write out what the actual issue is instead of mysteriously hinting at some potential problem?
Code like this "q, k, v = map(split_heads, tf.split(c, 3, axis=2))" will require me re-reading the paper section, then printing out all the tensors to think about what tensor would have which shape at which point. Instead of writing relatively useless linecomments like '#Transformer', I would comment all non-trivial shape modifications with the current layout, and what we are trying to achieve while modifying the layout.
The other issue of my original comment was not specifically on that codebase, but I am sure you would admit that the baselines code was pretty much exactly what I was writing about re: ml scripts. That's not to denigrate its incredible usefulness to the community.
Since you mentioned spinning up, I thought I would add a few comments on that as well:
I think the spinning up code base is good at making the code compact, and terrible at making sense of data flow for beginners. There are a lot of line comments, but they do not actually explain what is conceptually going on but often just repeat short-hands.
For example, look at the PPO implementation: https://github.com/openai/spinningup/blob/master/spinup/algo...
Here, the function is returning pi, logp, and log_p_pi (and v). Do you know how incredibly confusing the distinction between these is for beginners? In particular, there is no explanation why logp_pi even needs to be stored in the buffer.
We could recompute it from the states and stop the gradient when computing the likelihood ratio. A sensible tutorial-level comment here may be something along the lines of computing the likelihood in the same forward pass as computing the action, so we can later use it to compute the likelihood ratio. We could also later re-compute this from the buffered states.
I will stop here but I hope my point comes across, whenever I read code from your repos, there are some good parts (conciseness, cute numerical tricks) but there is as general missing sense of thoughtfulness on what the code is really trying to convey to a reader. It shows in the comments and it shows in the code organisation.
As a final note, I have seen this in many organisations and I do not mean to call you out. There is just this quality degradation that inevitably happens when nobody is incentivised (read: promoted, rewarded) to think about these things for an organisation.
Managers at all levels typically don't because they don't get close enough to the subtle issues on a day to day level. If you are lucky, you get senior individual contributors who still look at code and and raise the bar for the entire org. My genuine recommendation to you is to look for that, because a manager won't do that, and more fresh grads can't do it.
Very reasonable point that it is not clearly explained why you need to store logp_pi in the buffer. But the reason is that it would require additional code complexity to calculate it on the fly later. The likelihood ratio requires the denominator to be on the _old_ policy, so if you wanted to compute it on the fly, you would need to have a second policy in the computation graph to preserve the old policy while you change the current policy. You could not simply do a stop_gradient on the current policy and get the same results.
My personal feeling is that tutorial-style explanations like this don't fit nicely into code comment flow. As a result, most tutorial-style descriptions went into material on the Spinning Up website rather than into the code. It isn't 100% comprehensive, certainly, but RL has an enormous surface area (there are tons and tons of little details that teaching material could dive into) and I feel pretty good about what we were able to cover. :)
Writing a complex method in 174 lines is not elegant nor beautiful. Writing a well documented file that can take an engineer in a different specialty and bring them up to speed in 1,000 lines is.
Think about it like this: If you write a game engine, are you going to document each function with proofs that explain how the underlying Physics works for people who don't have such knowledge? No, you assume that people who read game engine code have read some basic physics books.
In practice, it's much harder to use that code, and we tend not to consume code like that internally. There's a real tradeoff!
Feel free to reach out by email, jachiam at openai.
It may be an indication that the whole thing needs refactoring though.
Many research scientist I've seen at large tech companies don't have good CS background and write code that's either unreadable, unscalable, unmaintainable, or buggy at times. They often have good ideas but don't do much software design.
Having been both at infrastructure teams and research teams, there are certain individuals in research orgs joining straight from school who think they are responsible with coming with a new and sexy thing and other engineers are responsible to run in production. It's like computing eigenvectors in Matlab or Python over a toy dataset and thinking you've done the bulk of the work of computing PageRank in production and should receive all the credit for a full search engine.
That attitude is a red flag to me.
As someone with experience in ML research I think this code is quite well written compared to what you typically see (a single function with hundreds of lines and dozens of if statements). I can immediately see what any of the functions does, and I haven't even read the paper.
def mult(a,b): return a*b
def mult(activation_a, activation_b): return activation_a*activation_b
Most of the code becomes really straightforward once you have. A lot of the magic constants are the result of multi page proofs (like the GELU constant) that would be impractical to put in the code.
Deep learning research really is a field that requires some amount of knowledge, and it's normal that you don't automatically understand state of the art code. Here is the GPT2 paper https://d4mucfpksywv.cloudfront.net/better-language-models/l...
I think this is because the "product" so to speak is often the papers themselves, not the code, but I'm not sure.
Complex graphs are literally a spaggeti of arrows, and this format actually is pretty readable (even though in pytorch it would be more readable). I guess they leave comments out because it's not really possible to understand each line on its own (unless it's an implementation detail); you have to read the paper to know what s going on
I remember reading a famous scientist (newton maybe) published a really accessible book on a subject, which was read by lots of lay persons and opened him up to lots of unwanted public attention.
So publishing in a more inscrutable way might be a way of assuring peer-to-peer communication.
Either that, or it's a labor of love where cleaning things up would detract from the forward momentum.
I'm skimming over some of the code at https://github.com/openai/gpt-2/blob/master/src/model.py and I can't help but feel frustrated at how unreadable this stuff is.
1. Why is it acceptable to have single-letter variable names everywhere?
2. There's little to almost no documentation in the code itself. It's unclear what the parameters of any given function mean.
3. There are magic constants everywhere.
4. Function names are so terse... ("gelu", "attn")