A New Inference Engine

For the past few months, I’ve been ded­i­cat­ing enor­mous swaths of my time to train­ing a small lan­guage model, work­ing from ba­sic prin­ci­ples. The goal is to pro­vide a com­plete over­haul of how Harper does its gram­mar check­ing.

In my first blog post about the sub­ject, I de­tailed some hy­po­thet­i­cal ideas for how we would im­ple­ment this change, but I did­n’t want to be too pre­scrip­tivist, so I did­n’t say any­thing con­crete. Now that some time has passed and I’ve earned a bet­ter un­der­stand­ing of the prob­lem do­main, I’m ready to share my vi­sion.

I’m writ­ing this out pub­licly for a new rea­sons:

• I see to doc­u­ment the de­ci­sion mak­ing process be­hind Harper’s ar­chi­tec­ture, so as to help new con­trib­u­tors un­der­stand how Harper works.
• Harper is de­signed to be ex­ten­si­ble. Some of the mod­els that we (Automattic) train will re­main pro­pri­etary, but if some­one in the com­mu­nity wishes to train their own for their own pur­poses (i.e., to sup­port a lan­guage we do not), they should be able to do so.

The Constraints

There are a num­ber of in­fer­ence en­gines that al­ready ex­ist. Why don’t we use or re­pur­pose one of those for Harper? In short: none of them work un­der the con­straints Harper is un­der.

Harper Runs On-Device

Most off-the-shelf in­fer­ence en­gines are de­signed for dat­a­cen­ter use. That’s fine — if you run your mod­els on spe­cial­ized hard­ware. Harper needs to run on any kind of de­vice, from the lat­est high-end MacBooks to ten year old ThinkPads. Not only that, but it needs to sip power. Our in­fer­ence en­gine needs to be able to slow it­self down (to sip power in the back­ground for non­crit­i­cal queries) and speed it­self up when the user de­mands Harper’s im­me­di­ate at­ten­tion.

Not All Our Models Work on Most Inference Engines

Harper will make use of an au­tore­gres­sive trans­former model for sen­tence rewrit­ing (more on that later), but that is not the only kind of model we will need. We will also be pro­vid­ing mod­els that per­form char­ac­ter mask­ing. What does that mean?

Instead of do­ing au­tore­gres­sive in­fer­ence (which can re­quire thou­sands of sep­a­rate for­ward passes through a model to pro­duce a sin­gle out­put), mask­ing mod­els do a sin­gle for­ward pass. They pro­vide char­ac­ter-level “masks” that we can use for an­other part of our gram­mar pipeline. For ex­am­ple, we could have a model (coupled with a LoRa, which we will talk about in a minute) per­form a sin­gle mask­ing pass over a user’s en­tire doc­u­ment, pro­duc­ing a list of lo­ca­tions they should have in­serted com­mas. From there, Harper can high­light those lo­ca­tions and of­fer to in­sert the rel­e­vant punc­tu­a­tion.

An au­tore­gres­sive model would have taken sev­eral min­utes to do this kind of analy­sis on con­sumer hard­ware. We can do it in frac­tions of a sec­ond.

Harper Runs Everywhere

On top of those two al­ready dif­fi­cult con­straints, Harper also needs to run any­where. We run na­tively on lap­tops via Harper Desktop, but we also run our en­gine in the ser­vice worker of browsers in WebAssembly. Whichever in­fer­ence en­gine we use must be able to run in all those places, us­ing the best avail­able ac­cel­er­a­tors.

LoRA Swap Speed Is Important

Different kinds of lint­ing tasks may need dif­fer­ent be­hav­ior from the same model. To ac­com­plish that while keep­ing model size tiny, we will NOT be train­ing any model for in­struc­tion fol­low­ing. In­stead, we will be dy­nam­i­cally load­ing and un­load­ing LoRAs from the model for mul­ti­ple in­fer­ence passes over the user’s doc­u­ment. One pass might an­a­lyze punc­tu­a­tion, an­other might an­a­lyze tone. To do this, we need to be able to swap be­tween mul­ti­ple LoRAs on the same model sev­eral times per sec­ond.

It’s not that most in­fer­ence en­gines can’t do this, but their in­ter­faces make it kludgey.

Andrew

Harper’s new in­fer­ence en­gine, which I am ten­ta­tively nam­ing “Andrew”, cur­rently re­sides in its own repos­i­tory. Soon, it will en­ter Harper’s monorepo. It’s based on Burn, a Rust li­brary for do­ing ML work any­where.

Andrew ful­fills all our re­quire­ments:

• It can run any­where (as you will soon see).
• It can load, un­load, and switch LoRAs in­stantly.
• Base model weights can be com­posed (like LEGO) to use the same weights when switch­ing be­tween batched mask­ing and au­tore­gres­sive in­fer­ence.
• It is op­ti­mized for con­sumer hard­ware.

Stay Tuned

Over the next few weeks, I’ll be re­leas­ing more in­for­ma­tion about how these mod­els are trained and how you can get ac­cess to them. In the mean­time, don’t be afraid to reach out if you have any ques­tions.