Google Search 淘气三千问:Q1~Q5

前言

两个多月前(2024年5月27日),Google 的一份名为 GoogleApi.ContentWarehouse 的 API 文档受到 SEO 圈的关注,由于这份文档的内容和 Google Search 副总裁 Pandu Nayak 在 2023 年美国司法部(DOJ)起诉 Google 的案件中的证词和 Google 的一些专利高度一致,因此其真实性被广泛认可。

后续有媒体称 Google 发言人回应了文档泄露的问题(没有承认、也没有否认):

A Google spokesperson sent us the following statement:

“We would caution against making inaccurate assumptions about Search based on out-of-context, outdated, or incomplete information. We’ve shared extensive information about how Search works and the types of factors that our systems weigh, while also working to protect the integrity of our results from manipulation.”

此前大部分讨论仅限于猜测 API 文档中的各种信号在排名算法中的作用,以及对谷歌是否在排名算法上欺骗了大家。很少人意识到,这篇 2500 页的文档可以作为以往 Google 公开论文的补充,一本叫做《如何构建一个世界级(成功的)搜索引擎》的武功秘籍撕下来的几页。

而偏偏我对这本武功秘籍非常好奇,试图在这几页上再加一些批注,就有了 《Google Search 淘气三千问》这个系列。这个系列会有几篇文章,我不知道,主要看我能想到多少问题。

为什么叫“淘气三千问”这个名字?可以把它看成是一种传承吧,懂的都懂,不懂的也不影响阅读。

为避免误导读者,这个系列所有的回答里,来自公开信息的我都会标注来源,没有标注来源的,你可以认为是 inaccurate assumptions about Search based on out-of-context, outdated, or incomplete information。当未来有了更新的信息,我有可能回到博客来更新这些 assumptions(公众号文章无法更新)。

如果你有新信息可以给我,或者纠正文中的错误,欢迎评论或者到公众号“边际效应”私信,谢谢!

Q1: Google 的索引分了几层(Tier)?依据什么?

Google 在 2012 年的论文《Indexing the World Wide Web: The Journey So Far 》中提到产业实践中大规模索引都是会分成多个桶(tier),一般按照更新频率来分:

The way we have described search indices so far makes a huge assumption: there will be a single unified index of the entire web. If this assumption was to be held, every single time we re-crawled and re-indexed a small set of fast-changing pages, we would have to re-compress every posting list for the web and push out a new web index. Re-compressing the entire index is not only time consuming, it is downright wasteful. Why can we not have multiple indices -- bucketed by rate of refreshing? We can and that is what is standard industry practice. Three commonly used buckets are:

1. The large, rarely-refreshing pages index
2. The small, ever-refreshing pages index
3. The dynamic real-time/news pages index

...

Another feature that can be built into such a multi-tiered index structure is a waterfall approach. Pages discovered in one tier can be passed down to the next tier over time.

在 Google ContentWarehouse API 里有这样一段 :

GoogleApi.ContentWarehouse.V1.Model.PerDocData

* scaledSelectionTierRank (type: integer(), default: nil) - Selection tier rank is a language normalized score ranging from 0-32767 over the serving tier (Base, Zeppelins, Landfills) for this document. This is converted back to fractional position within the index tier by scaled_selection_tier_rank/32767.

可以看到,Google 仍然是把索引分了 3 层,现在我们有了它们的名字,分别是:Base(基础)、Zeppelins(飞艇) 和 Landfills(垃圾填埋场)。在每一层之内,scaledSelectionTierRank 这一归一化分数决定了它所在位置的分位数。分位数最大值是 32767,猜测也许是 Google 在索引存储里只给它留了 15 bits(2^15=32768)。

但从索引分层的名字来看,这三层并不(全)是按照时效性分的,至少第三层,看着是按照文档质量分的。因为你把文档放到“垃圾填埋场”中,大概率因为它的质量较差而不是不再更新。那么 scaledSelectionTierRank 也许就代表了层内的文档质量等级。

网友 avanua 对这三层的命名提供了一个解读,我觉得非常合理,因为我一直困惑第二层为什么叫做 Zeppelins:

我觉得 Tier 命名和质量无关,可能只是用来描绘更新频率:
​Zeppelins 在气流中起起伏伏
Base
​Landfills 几乎不会再翻动,上下层叠关系是固定的

Q2: Tier 内的 scaledSelectionTierRank 有什么作用?

在 《Indexing the World Wide Web: The Journey So Far》中提到,在倒排拉链中最好按照文档实际的影响力对文档列表进行排序。如果仅仅是这样,那么只需要知道文档 0 比文档 10000 更重要即可,那么额外记录一个打分的目的,其实是可以让这个分数参与排序过程。文档在某个 Query 下的得分,是文档影响力得分乘以文档在 Query-term 下的权重。

Since it made sense to order the posting lists by decreasing term frequency, it makes even more sense to order them by their actual impact. Then all that remains is to multiply each posting value by the respective query term weight, and then rank the documents. Storing pre-computed floating-point document scores is not a good idea, however, since they cannot be compressed as well as integers. Also, unlike repeated frequencies, we can no longer cluster exact scores together. In order to retain compression, the impact scores are quantized instead, storing one of a small number of distinct values in the index.

从上文中有理由认为,scaledSelectionTierRank 就是文中提到的量化以后的文档影响力得分,量化就是将其归一化到 32768 个分档之中。

Q3: Google 搜索系统主要分成几个部分?

通过 API 和其它公开文档,目前我能够分析出来的搜索系统组成部分有以下这些。随着阅读的深入,可能还会有新的部分加进来。

爬虫系统:Trawler

在 Google ContentWarehouse API 中有一系列 API 以 Trawler 为前缀,并且从上下文中看出来 Trawler 是一个实体系统并且有一个研发团队。

GoogleApi.ContentWarehouse.V1.Model.TrawlerCrawlTimes
GoogleApi.ContentWarehouse.V1.Model.TrawlerFetchReplyData
GoogleApi.ContentWarehouse.V1.Model.TrawlerHostBucketData

* TotalCapacityQps (type: number(), default: nil) - The following four fields attempt to make things simpler for clients to estimate available capacity. They are not populated yet as of 2013/08/21. Even after they are populated, they may change. So talk to trawler-dev@ before you use the fields. Total qps for this hostid

去重系统:WebMirror

在 Google ContentWarehouse API 里有这样一段:

GoogleApi.ContentWarehouse.V1.Model.CompositeDocAlternateName

Alternate names are some urls that we would like to associate with documents in addition to canonicals. Sometimes we may want to serve these alternatenames instead of canonicals. Alternames in CompositeDoc should come from WebMirror pipeline.

每个 CompositeDoc 都有一些替代的 URL,这些 URL 来自 WebMirror 流水线,那么 WebMirror 应该是识别重复文档的一套系统。

离线索引构建系统:Segindexer + Alexandria

在 Google ContentWarehouse API 里有这样一段:

GoogleApi.ContentWarehouse.V1.Model.AnchorsAnchor

* sourceType (type: integer(), default: nil) - ... In the docjoins built by the indexing pipeline (Alexandria), ...

所以 Alexandria 应该是建库系统。而 Segindexer 和 Alexandria 曾经并行出现过:

GoogleApi.ContentWarehouse.V1.Model.ClassifierPornClassifierData

* imageBasedDetectionDone (type: boolean(), default: nil) - Records whether the image linker is run already. This is only used for Alexandria but NOT for Segindexer.

考虑到关键的表示原始文档内容的 compositedoc.proto 在 Segindexer 目录下:

GoogleApi.ContentWarehouse.V1.Model.NlpSaftDocument

* bylineDate (type: String.t, default: nil) - Document's byline date, if available: this is the date that will be shown in the snippets in web search results. It is stored as the number of seconds since epoch. See segindexer/compositedoc.proto

从名字和上述信息有理由怀疑 Segindexer 是在 Alexandria 之前,决定了索引分层,或者分 vertical 的一个分类模块。

在线索引服务系统:Mustang 和 TeraGoogle

在 《Indexing the World Wide Web: The Journey So Far》中我们知道,TeraGoogle 是 Google 在 2005 年实现的一套 large disk-based index 服务系统。而在 Google ContentWarehouse API 里有这样一段:

GoogleApi.ContentWarehouse.V1.Model.CompressedQualitySignals

A message containing per doc signals that are compressed and included in Mustang and TeraGoogle.

这里将 Mustang 和 TeraGoogle 并列,有理由认为 Mustang 是 2005 年之后 Google 开发的一套替代或者部分替代 TeraGoogle 的在线索引服务系统。

实时索引服务系统:Hivemind(muppet)

Realtime Boost 这篇文档中有一张图,揭示了一个新网页进入索引的过程,其中实时索引的系统叫做 Hivemind。但 Hivemind 应该是个集群名,实际的索引系统方案可能是叫做 muppet。因为在 RealtimeBoost Events - DesignDoc 这篇文档中也提到了一个用 muppet 支持的系统,叫做 ModelT,这是一个索引实时事件的索引系统。

实时索引的建库系统可能与非实时索引不同,叫做 Union。

查询汇聚系统:SuperRoot

在 Google ContentWarehouse API 中多次出现 SuperRoot 这一模块,而在 Jeff Dean 2009 年 WSDM 的 《Challenges in Building Large-Scale Information Retrieval Systems》 分享第 64 页,SuperRoot 被描述为聚合 Web、Images、Local、News、Video、Blogs 和 Books 所有检索子系统的汇聚模块,这个定位也许没有变。

Query 改写模块:QRewrite

RealtimeBoost Events - DesignDoc 这篇文档中说:

Currently in production RealtimeBoostServlet runs in QRewrite detects spikes on news documents published that match the SQuery (including syns). It does so by issuing one or two RPCs to Realtime-Hivemind.

The RealtimeBoostResponse containing the Spike is sent down to Superroot in the QRewrite response and it is currently used by a few Search features (such as TopStories) to trigger faster and rank fresher documents for queries that are spiking for a given news event.

所以 QRewrite 是会被 SuperRoot 调用的一个 Query 改写模块,它会将用户实际发起的搜索 Query 改写成 SQuery,S 可能是代表 Superroot。而到真正发起检索时,会进一步转成检索 Query,比如发给 muppet 的 Query 叫做 MQuery。

摘要模块:SnippetBrain

在 Google ContentWarehouse API 里有这样一段:

GoogleApi.ContentWarehouse.V1.Model.MustangReposWwwSnippetsSnippetsRanklabFeatures

* displaySnippet (type: GoogleApi.ContentWarehouse.V1.Model.QualityPreviewRanklabSnippet.t, default: nil) - Snippet features for the final chosen snippet. This field is firstly populated by Muppet, and then overwriten by Superroot if SnippetBrain is triggered.
  • 看起来 SnippetBrain 是一个可选的摘要生成模块。

入口服务:GWS

Google Web Server,这个大家都知道,还有 Wikipedia 词条

Q4: TeraGoogle 是怎样一套系统?

根据论文《Indexing the World Wide Web: The Journey So Far》和专利《US7536408B2: Phrase-based indexing in an information retrieval system》,TeraGoogle 应该有以下几个属性:

  • Disk-based Index:索引存储在磁盘上,在需要的时候读入到内存中,而且往往不需要全部读入,针对重要的文档有一些优化;
  • Phrase-based Indexing:构建索引的时候不仅仅有 term 索引,还会建设多 term 的短语索引,这样索引库里会有更多的倒排链;
  • Document-Partitioned Index:将索引分 Shard 的时候,按照文档进行分片,即同一批文档的所有拉链放在同一个 Shard 上,这样每个 Shard 上有所有的拉链,查询在一个节点内即可完成。在论文中只对比了 Document-Partitioned 和 Term-Partitioned 二者的差异,在 Jeff Dean 2009 年 WSDM 的 《Challenges in Building Large-Scale Information Retrieval Systems》 分享第 17 页确认了 Google 的选择。

Q5: Google 的文档是什么概念?

从 Google ContentWarehouse API 里:

GoogleApi.ContentWarehouse.V1.Model.CompositeDoc
Protocol record used for collecting together all information about a document. 

可以看到 CompositeDoc 是在系统里非常重要的概念,它定义了一个文档的所有信息。在它的所有字段中我们发现,url 又是一个可选的字段,这也就是说,文档并不一定需要是一个网页。像 localinfo,看起来就像是一个 POI 信息。也就是说在 Google 的系统里,不一定只有网页索引,可能每个 POI 点、图片、商品也是一种文档,所以它使用 CompositeDoc (复合文档) 而不是 WebPage 作为整个系统里对文档的刻画。

估值最高的 AI 搜索独角兽 Perplexity 使用倒排索引做 RAG

3 月份我曾写了一篇博客《在 LLM 时代我们是否还需要倒排索引? 》,探讨了在向量数据库崛起的情况下倒排索引仍然存在的价值。

也许仍然有人对传统搜索技术弃如敝履,但前几天我正好看到一个 Lex Fridman 对 Perplexity CEO Aravind Srinivas 的采访,他表示 Perplexity 仍然使用(可能不是仅使用)倒排索引和 BM25 传统算法进行全网搜索。作为估值 30 亿美元的 AI 搜索独角兽公司,在全网搜索上却也选择了倒排索引,而且访谈中非常强调倒排的重要性,这应该算是印证了我的观点吧。

他的这些思考对我也有一些启发,为了保持互联网的记忆,我将这些相关节选的原文和翻译记录如下:

Aravind Srinivas 2024年6月20日采访节选

这个 Podcast 有 3 个小时,可以从 2:03:52 开始看。

Lex Fridman (02:03:52)

这是一个通过(将内容)嵌入到某种向量空间实现的完全的机器学习系统吗?
Is that a fully machine learning system with embedding into some kind of vector space?

Aravind Srinivas (02:03:57)

它并不是纯粹的向量空间。并不是说内容一旦被获取,就会有某种 BERT 模型在所有内容上运行并将其放入一个巨大的向量数据库中进行检索。
It’s not purely vector space. It’s not like once the content is fetched, there is some BERT model that runs on all of it and puts it into a big, gigantic vector database which you retrieve from.

这并不是那样的,因为将网页的所有知识打包到一个向量空间表示中是非常困难的。首先,向量嵌入在处理文本时并没有想象中的那么神奇。理解一个文档是否与一个特定检索词相关非常困难。它应该是与检索词中的人物有关,还是与检索词中的特定事件有关,更或者它是基于对检索词更深层次的理解,使得检索词可以应用到另外一个人物,他也需要被检索到?什么才是(向量)表示应该捕捉的内容?大家可能会不断争论。而让这些向量嵌入具有不同维度,彼此解耦并捕捉不同语义是非常困难的。
It’s not like that, because packing all the knowledge about a webpage into one vector space representation is very, very difficult. First of all, vector embeddings are not magically working for text. It’s very hard to understand what’s a relevant document to a particular query. Should it be about the individual in the query or should it be about the specific event in the query or should it be at a deeper level about the meaning of that query, such that the same meaning applying to a different individual should also be retrieved? You can keep arguing. What should a representation really capture? And it’s very hard to make these vector embeddings have different dimensions, be disentangled from each other, and capturing different semantics.

顺便说一句,这只是排序部分的内容。还有索引部分——假设你有一些处理后的 URL,排序模块会基于你提问的检索词,从索引中召回相关的文档和一些打分。这就是为什么,当你的索引中有数十亿个页面,而你只需要前 K 个结果时,你必须依赖近似算法来获取这些前 K 个结果。
This is the ranking part, by the way. There’s the indexing part, assuming you have a post-process version for URL, and then there’s a ranking part that, depending on the query you ask, fetches the relevant documents from the index and some kind of score. And that’s where, when you have billions of pages in your index and you only want the top K, you have to rely on approximate algorithms to get you the top K.

Lex Fridman (02:05:25)

所以这就是排序,但是将页面转换为可以存储在向量数据库中的内容,这一步似乎非常困难。
So that’s the ranking, but that step of converting a page into something that could be stored in a vector database, it just seems really difficult.

Aravind Srinivas (02:05:38)

它并不必须全存在向量数据库中。你可以使用其他数据结构和其他形式的传统检索。有一种算法叫做 BM25,它正是为此设计的,是 TF-IDF 的一个更复杂的版本。TF-IDF 是词频乘以逆文档频率,是一种非常老派的信息检索系统,实际上即使在今天也仍然非常有效。BM25 是它的一个更复杂的版本,仍然在许多排名中击败大多数嵌入。当 OpenAI 发布他们的嵌入时,围绕它有一些争议,因为在许多检索基准测试中,它甚至没有击败 BM25,这并不是因为他们做得不好,而是因为 BM25 太好了。这就是为什么仅靠纯粹的嵌入和向量空间并不能解决搜索问题。你需要传统的基于词项的检索,你需要某种基于 NGram 的检索。
It doesn’t always have to be stored entirely in vector databases. There are other data structures you can use and other forms of traditional retrieval that you can use. There is an algorithm called BM25 precisely for this, which is a more sophisticated version of TF-IDF. TF-IDF is term frequency times inverse document frequency, a very old-school information retrieval system that just works actually really well even today. And BM25 is a more sophisticated version of that, that is still beating most embeddings on ranking. When OpenAI released their embeddings, there was some controversy around it because it wasn’t even beating BM25 on many retrieval benchmarks, not because they didn’t do a good job. BM25 is so good. So this is why just pure embeddings and vector spaces are not going to solve the search problem. You need the traditional term-based retrieval. You need some kind of NGram-based retrieval.

Aravind Srinivas 2023年12月14日采访节选

其实在去年 Unsupervised Learning 的 Jacob Ephron 采访 Aravind Srinivas 时,他也有过类似的表述,但是没有像最新一次采访那样强调倒排的重要性。这个 Podcast 也很长,可以从 28:08 开始看。

Aravind Srinivas (28:08)

很多人认为,既然我们在网页搜索的 RAG 方面如此擅长,那么 Perplexity 就能轻松搞定内部搜索。不,这是完全不同的两回事。
A lot of people think that because we are so good at RAG for web search, Perplexity will just nail internal search. No, it's two completely different entities.

Google Drive 的搜索之所以糟糕是有原因的。你可能会想,Google作为网页搜索的王者,怎么会这么差劲?他们之所以差劲,是因为索引机制完全不同,需要训练的嵌入模型也完全不同。这不仅仅是嵌入模型的问题,还包括你如何对页面进行摘要、如何进行文本召回——如何使用传统的基于TF-IDF的倒排索引来构建弹性索引,这些都大不相同。
There is a reason why search on Google Drive sucks. Like would you expect Google the king of web search to be so bad? They're bad because of a reason that the indexing is so different, the embeddings that you got to train are so different. Not just the embeddings, but even the way you snippet a page, your text retrieval——the elastic index that you're building with traditional TF-IDF based inverted indexes are so different.

所以需要某家公司专注于这种场景,就像我们专注于网页搜索场景一样。RAG 是一项非常艰巨的任务,在生成式 AI 之外还有很多工作需要完成。这不仅仅是训练一个大型的嵌入模型就能解决的问题。
That you need a company to just obsessively focus on that use case. Just like how we are obsessively focused on the web search use case. So RAG is going to be pretty hard and there's a lot of work that needs to be done outside of generative AI. It's not just training a large embedding model and you're done.

我记得当 OpenAI 发布嵌入 API 时,Sam Altman(OpenAI的CEO)在推特上说,下一个万亿公司可能只需要接入这个 API 就可以建立起来。这虽然是一种很好的营销方式,但事实并非如此。他当时说的是一万亿美元。所以,当听到有人声称他们已经解决了 RAG 问题时,我会非常谨慎。他们可能只是在某个场景下做得很好。
I remember like when OpenAI releases embedding API, Sam Altman tweeted the next 1 trillion company can just plug into this API and be built. That's not true. It's a good way to market it. But that's not true at all. Sorry, he said 1 trillion dollars. So I think that's why I would be very careful when somebody makes claims that they've solve RAG. They probably can do it really well for one use case.

此外,在排名中还有很多其他因素需要考虑,才能使答案真正出色。即使 LLM 最终决定了哪些链接用于答案,你也不能仅仅将一些垃圾信息放入 prompt 中,就指望它能神奇地只选择最相关的链接,并在答案中给出引用。
You know and also there are so many more things to handle in the ranking. That'll make the answer really good. Cuz even though the LLMs are finally the ones that are taking which links to use for the answer, it's not like you can just dump garbage into the prompt and it'll just be magically so good that it'll only take the top most relevant links in the answer and give you the citations with them.

事实上,你向这些长上下文模型提供的信息越多,最终出现幻觉的可能性就越大。因此,你实际上需要在检索模块上投入大量工作,不仅仅是嵌入向量,在索引、嵌入向量和排序上都要投入。排序也应该包含很多除了向量内积之外的其他信号,具体是哪些信号事实上取决于你的场景。
In fact the more information you throw at these really long context models, the more chances that you have a hallucination at the end. So you actually have to do a lot of work in the retrieval component, not just the embeddings, the indexing, the embeddings and the ranking. Ranking should also have a lot of signals outside of just the vector dot products. And then what is those signals are really depend on your use case.

分析(在2024年8月)

从以上不同时期 Aravind 的公开披露的信息分析,几乎可以说 Perplexity 在当前时间点,在召回阶段主要(如果不是全部的话)依赖倒排索引,在排序阶段会用到嵌入向量和其它信号,并且他们很重视除了嵌入向量之外的其它信号。

LLM 推理优化 Prefix Caching 及其实现

Prefix Caching

上文中提到,Prompt 计算(Prefill 阶段)和生成阶段的计算特性很不相同。为避免重复计算,所有框架的 prefill 阶段主要作用就是给迭代的生成阶段准备 KV Cache。但这些 KV Cache 仅仅是为单次生成式请求服务的,那很自然的一种想法就是,KV Cache 能不能跨请求复用?

在某些场景下,多次请求的 Prompt 可能会共享同一个前缀(Prefix),比如拟人 Agent 的人物设定,文档阅读理解时的文档内容等。这些情况下,很多请求的前缀的 KV Cache 计算的结果是相同的,像一般互联网服务的请求缓存一样,可以被缓存起来,给下一个请求复用。

限制

但 KV Cache 跟其它服务缓存不一样的地方是,它太大了,以至于(目前)很难通过 Redis/Memcache 这种分布式缓存服务存取。比如对 13B LLM 模型来说,在 FP16 精度下单 token 的 KV Cache 大约是 1MB,假设要缓存的前缀有 500 个 token(大约800多个汉字),那就是 500MB。一般来说,我们不会每次请求去从分布式系统里读取/传输 500MB 的缓存,甚至都不会每次请求从内存往显存中拷贝 500MB 的缓存,所以大部分情况下,prefix cache 都会放在显存里。

这也就意味着,如果你想命中 prefix cache,必须把相同 prefix 的请求发到同一张 GPU卡上才行。

实现

由于不是普遍需求,加上前面说的限制,prefix caching 作为一个加速特性,不是很受关注,一般也不是默认开启的。各框架的实现和配置略有差异,这里简单做下记录,便于回顾。

刚开始 vLLM 的实现是给 generate 接口增加一个 prefix_pos 参数,通过 prefix_pos 输入参数为每个请求指定 prefix 长度,为 prefix 建一个带淘汰的哈希缓存。后来觉得这样做使用上不够便利,升级成了自动前缀缓存,即将 prompt 的 kv cache 分成 block,然后为 block 建设 LRU 缓存机制,这样就不必在接口上使用 prefix_pos 指定哪部分是 prefix 了。自动前缀缓存功能默认是不开启的,开启的配置项为 --enable-prefix-caching

TensorRT-LLM 与 vLLM 后来的实现类似,也是实现了 block kv cache,配置项是 enableBlockReuse,默认也是不开启的。代码未开源,无法看到实现。

Lmdeploy 的 PythonTurboMind C++ 版本的 prefix caching 功能都已经有了 PR,但现在(20240425)看还没有合入主干。有意思的是它没有使用 hash block 对应 token_id 子串的所有 token_id 前缀然后组成哈希表的方式,而是用 hash 当前 block 对应的 token_id 子串然后组成 trie 树的缓存管理结构。默认的参数名与 vLLM 相同,也叫做 --enable-prefix-caching。

HuggingFace TGI 现在看起来还没实现 Prefix Caching 功能。

Prompt Caching

除了 Prefix Caching 这种比较直观的工程优化,现在也有一些研究在看 Prompt 的其它缓存机制。比如设计一种机制让 prompt 模块化,不仅可以复用 prefix,还能复用中间的部分;或者通过 query 的相似性复用其它 query 的 prompt

但目前看实现上都过于复杂,比如第一种要求模型使用不连续的 poition_id,这样就可以插入 token,但这种方式对 attention 的计算机制有一定的影响,难以说明它对效果的影响。

LLM 推理优化 Continuous Batching 及其实现

原理

Continuous Batching 是 LLM 推理优化的一项技术,作为这篇文章的知识背景不再赘述,目前流传最广的参考资料是这篇:《How continuous batching enables 23x throughput in LLM inference while reducing p50 latency》。它也有中文翻译,感兴趣可以搜一下,先看看。

图片来自:Anyscale 博客

虽然这篇资料介绍了它的主要原理,尤其是这张简单易懂的图,但是实现与原理是存在差异的,因为工程实现要解决很多现实问题。

(Re)scheduling 开销

如原文所说,Continuous batching 还有个别名,叫做:batching with iteration-level scheduling,这里的 iteration 就是指一次 decode 计算。也就是说在每次 decode 的迭代过程中,做 batch 的调度调整。

但调度本身不是无代价的,它可能涉及到接收和处理新的输入请求,重新组织输入数据的形状,甚至各种状态的重新初始化,这些都需要消耗 CPU 时间。这也就意味着在这段时间里,GPU 是闲着的,GPU 没有得到充分利用。

所以在实现时,程序并不会真的每个 iteration 都做 scheduling,目前看到有两种做法:

  • 合理间隔调度。比如每 16 次 decode 计算后,检查一下是否有新的输入,以及是否有空闲的槽位,然后对 batch 做一次调度调整。这能够显著降低调度的开销(TGIlmdeployvLLM)。
  • 排队比例调度。比如当前 batch 中有 10 个请求的 decode 正在进行,而排队中有 12 个请求,超过了排队比例 1.2,那么就启动一次调度调整(TGI)。

KV Cache 重读

如果真的像图中那样,每个生成 Token 的 decode iteration 与一个 prompt token 的计算对齐,那 KV Cache 的利用就会很糟糕。因为它们需要在 Global Memory 与 Shared Memory 之间反复搬运,而且每次搬进来以后只被使用一次。

这本质上是 prefill 阶段(prompt 计算)与生成阶段的计算特性不同导致的,prefill 阶段并不适合 one-by-one 的 token 计算。

所以在实现时,程序并不会真的做 prefill 和生成的 token 对齐调度。目前看到的调度方法有三种:

  • 在重调度时,如果有新的请求进来,那么将新请求的 prefill 计算和要进行的 decode 计算做合并(Orca、vLLM-prev)。
  • 在重调度时,如果有新的请求进来,那么先对新的请求做 prefill 计算,然后再合并所有进行中的请求做 decode 计算(TGI、vLLM、lmdeploy)。
  • 先根据 decode 耗时估算出来每次 decode 同时能做 prefill 的 token 数量,在重调度时,如果有新的请求进来,对新请求的 prefill 按照上面的估算进行分块,然后将分块的 prefill 和其它请求的 decode 计算融合在一起,一定程度上减少了 KV Cache 重读的次数,又避免先做 prefill 计算带来的 Token 生成延时增加(Sarathi-Serve+vLLM)。

可调优能力

LLM 推理服务往往用于生产环境,而生产环境面临的情况是复杂多样的。

  • 对于做阅读理解的应用来说,Prompt 可能会非常长,但生成的内容可能会非常短,开发人员可能会更追求吞吐;
  • 对于聊天应用来说,Prompt 可能较短,生成的内容也不会太长,开发人员可能会更追求延迟;
  • 对于创作类应用来说,Prompt 可能很短,生成的内容会更长,开发人员可能会更追求首 Token 延迟。

对 Continuous Batching 实现来说,就要求它调度策略尽量清晰,并且参数可调。所以更灵活的实现,未来可能会更受欢迎。

Logits of API-Protected LLMs Leak Proprietary Information

看到一篇挺有意思的论文,大开脑洞,没想到还能这么玩,做一下粗读的笔记。

论文

标题:《Logits of API-Protected LLMs Leak Proprietary Information》,链接: https://arxiv.org/pdf/2403.09539.pdf

假设条件

典型 LLM 需要将最后一个 Transformer 块输出的嵌入向量转成要输出的 Token,这一步往往通过一个线性变换加 softmax 获取那个最大概率的 tokenid。

线性变换的权重是一个 vocabulary size * hidden size 的矩阵,比如 llama2-7B 的词表大小是 32000,hidden size 是 4096,那么线性变换权重矩阵的尺寸就是 32000x4096。这个矩阵再与 4096x1 的嵌入向量相乘,得到的就是 32000x1 的 logits 向量,其中每一个元素对应着一个词表中的 token 作为最终输出的概率。

上面这只是假设,也许 GPT 使用的是一个非线性变换,那论文内容可能就不成立了。

数学原理

这个线性变换将一个 4096 维的向量映射到了一个 32000 维的向量,从线性代数的角度来看,这是一个低维向高维的映射,所以它肯定不是一个满射(onto mapping)。也就是说,这个映射的像空间(image)只是 32000 维实数空间的一个子空间(subspace),而且这个像空间的秩(rank)最多是 4096。

这意味着可以找到不多于 4096 个线性无关的基向量(basis),使得这个像空间的每一个元素都能表示为这些基向量的线性组合。假设能采集到 4096 个线性无关的输出 logits,那这些 logits 就构成了像空间的一组基向量。

反过来想,如果你不知道 LLM 的 hidden size,那么你可以通过采集足够多的输出 logits,以保证有足够多的线性无关的向量。然后对矩阵进行奇异值分解(singular value decomposition),可以通过非 0 的奇异值个数推导出矩阵的秩。这个秩应该接近于模型的 hidden size。

逆向恢复 logits

遗憾的是,很多模型的 API 并没有输出完整的 logits 矩阵,但幸运的是,OpenAI 的 API 支持输出最多 top 5 个 token 的 logprobs,并且支持 logit_bias 干预 token 的输出。那就给了反复通过 API 调用来逆向恢复 logits 向量的可能。

但是具体方法我没看,粗读嘛,知道能做到就行了,有用到的时候再看吧。还有另一篇文章《Stealing Part of a Production Language Model》分析了在没有 logit_bias 甚至没有 logprobs 时该如何恢复 logits,我也没看,记录下链接 https://arxiv.org/pdf/2403.06634.pdf

无法输出的 Token

这篇论文还介绍了很多其它应用,太长没有看。比较有意思的一个引用是,在将嵌入向量映射到 logits 的过程中,如果一个 token 的嵌入向量在其它 token 的嵌入向量组成的凸包的内部,它就永远不可能被输出。扫了一眼引用的论文,证明没看懂,大致意思是 softmax 权重矩阵的低秩特性导致了可能输出 token 的排列在线性变换后不会出现在子空间里?实话说我感觉不像是很严谨的数学证明。。。

在 LLM 时代我们是否还需要倒排索引?

近些年,EBR(基于文本嵌入向量的召回)的强大语义召回能力让基于属性索引的传统倒排索引技术黯然失色,即使对专业搜索引擎来说,EBR 的应用也是越来越广泛 [1,2,3] 。尤其在 LLM(大语言模型)激起的 RAG(检索增强生成)技术框架下,大部分人似乎已经忘记了倒排索引,向量数据库成为 RAG 的标配。

但在 LLM 时代倒排索引真的没有用武之地了吗?我尝试列一下自己的思考。

  • Embedding 向量缺少 ground truth,但是倒排有。你无法通过直接观察或者测量,来明确一个向量指代的一段文本是否一定包含某些信息。但是如果这段文本在某个 term 的倒排拉链里,你可以从一定程度上明确这段文本包含了一些相关信息。
  • term 命中也是一种 attention。在训练模型时,我们总是希望 LLM 能关注到 context 中应该关注的信息,而忽略其它无关内容。那跟用户问题或者指令中某些 term 相关的内容,应该需要更多的关注。其实也可以类比人类肉眼的信息查找,人们也总是会扫到某些关键词,然后再仔细阅读它的上下文篇章。
  • 不基于倒排做召回,仍可以用倒排做粗筛。倒排作为一种可以查询 term 命中的高效结构,对 EBR 也许可以起到补充作用。例如对于某些 EBR 效果不够理想,误召回概率较高的场景下,对得分比较低的文档用命中信息作一次粗筛,能显著提升送给模型的 context 质量,也能减少对 LLM 计算资源的浪费。
  • term 命中的位置信息和权重不再重要。对于 LLM 来说,它会自行关注到 context 中需要关注的信息,不再需要位置信息或者权重来指示文本中哪些部分更重要。也就是说,倒排只需要解答 term 在文本中是否出现的问题,而不需要回答出现几次、在哪里出现的问题。
  • 也许倒排不再用 term,而是 token。term 依赖于切词,有一定的语义含义,term 集合空间一般有百万甚至千万的量级。但现在 LLM 大部分使用 BPE(Byte-Pair Encoding)分词器,token 集合空间只有几万到十几万的量级。使用 token 将显著减少倒排链的数量而优化其性能,但 token 存在没有归一化、分词边界不对齐的问题,是否可用还有待验证。

参考

[1] Guo, Ruiqi, et al. "Accelerating large-scale inference with anisotropic vector quantization." International Conference on Machine Learning. PMLR, 2020

[2] Huang, Jui-Ting, et al. "Embedding-based retrieval in facebook search." Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining. 2020.

[3] Liu, Yiding, et al. "Pre-trained language model for web-scale retrieval in baidu search." Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining. 2021.

配置 dcgmi 遇到的问题

dcgmi 是 Nvidia datacenter-gpu-manager 的命令行程序,可以用来采集 GPU 各类子资源的利用率数据,揭示的数据比 nvidia-smi 更详细,也更便于对接监控系统(比如 Prometheus)。这次我主要想用它来看模型训练过程中的 NVLink 带宽使用情况。

在一个已经完成 Nvidia 训练环境配置的服务器上安装 dcgmi,主要有以下几步:

1. 如果服务器上有 NVSwitch,比如 A100/800,或者部分 V100,需要先安装 nvidia-fabric-manager,注意 nvidia-fabric-manager 的版本要和 GPU 驱动版本严格一致,包括小版本号(nvidia-smi 显示的驱动版本)。

2. 如果服务器上有 NVSwitch,还需要安装 libnvidia-nscq,版本也要和驱动严格一致。

3. 安装 datacenter-gpu-manager

4. 启用并启动 nvidia-fabricmanager 服务,注意服务名(nvidia-fabricmanager)、包名(nvidia-fabric-manager)和进程名(nv-fabricmanager)的区别。根据启动成功与否,可检查 /var/log/fabricmanager.log 中的错误信息。

systemctl enable nvidia-fabricmanager
systemctl start nvidia-fabricmanager

5. 启用并启动 nvidia-dcgm服务,注意包名(datacenter-gpu-manager)、服务名(nvidia-dcgm)和实际命令/进程名(nv-hostengine)的区别。根据启动成功与否,可检查 /var/log/nv-hostengine.log 中的错误信息。

systemctl enable nvidia-dcgm
systemctl start nvidia-dcgm

但是整套配置下来,我还是遇到了不少问题,有些是配置问题,有些是文档未明确说明的问题。甚至我找到 Nvidia 的售后工程师咨询,他们都没法给我答案,后来还是我自己摸索出来了,做个记录。

1. nvidia-dcgm 服务启动成功,但又没完全成功

启动服务时候没有报错,但是执行监控命令时出错,其实启动的时候 /var/log/nv-hostengine.log 就已经出现无法加载 nscq 库的错误,只不过没影响服务的启动。因为 nscq 是对 nvswitch 的查询,用户如果不需要查询 nvswitch 的数据,可以接受这个错误。

ERROR [13862:13862] [[NvSwitch]] Could not load NSCQ. dlwrap_attach ret: Can not access a needed shared library (-79): If this system has NvSwitches, please ensure that the package libnvidia-nscq is installed on your system and that the service user has permissions to access it. [/workspaces/dcgm-rel_dcgm_3_3-postmerge/modules/nvswitch/DcgmNvSwitchManager.cpp:798] [DcgmNs::DcgmNvSwitchManager::AttachToNscq]
ERROR [13862:13862] [[NvSwitch]] AttachToNscq() returned -25 [/workspaces/dcgm-rel_dcgm_3_3-postmerge/modules/nvswitch/DcgmNvSwitchManager.cpp:632] [DcgmNs::DcgmNvSwitchManager::Init]
ERROR [13862:13862] [[NvSwitch]] Could not initialize switch manager. Ret: DCGM library could not be found [/workspaces/dcgm-rel_dcgm_3_3-postmerge/modules/nvswitch/DcgmModuleNvSwitch.cpp:34] [DcgmNs::DcgmModuleNvSwitch::DcgmModuleNvSwitch]

虽然我已经安装了 libnvidia-nscq 库,但 nv-hostengine 就是找不到 NSCQ 对应 so。我在这个问题上困扰了很久,Nvidia 的工程师也没法给我解释原因,只是一遍一遍地让我去看文档,甚至重装系统。

我觉得他的建议不靠谱,后来还是自己找到了原因,那就是 datacenter-gpu-managerlibnvidia-nscq 也有隐式的版本依赖。虽然在文档和包依赖中完全没有体现这种依赖关系,但我通过降级 datacenter-gpu-manager 到与 libnvidia-nscq 时间上更接近的版本解决了这个找不到 so 的问题。

之所以想到去尝试降级版本,还是因为 Nvidia 各种环境和驱动版本的强绑定,让我免不了去怀疑一下这个。

2. 查询部分指标 DCGM_FI_PROF_* 时出错

在执行 dcgmi dmon -e 449,1011,1012 时,命令行显示如下错误:

Error setting watches. Result: The third-party Profiling module returned an unrecoverable error

查看 /var/log/nv-hostengine.log,有如下错误日志:

ERROR [225876:233276] [[Profiling]] [ProfModule][PerfWorks] Got status 1 from NVPW_DCGM_PeriodicSampler_BeginSession() on deviceIndex 0 [/workspaces/dcgm-rel_dcgm_2_3-post
merge/dcgm_private/modules/profiling/DcgmLopGpu.cpp:351] [DcgmLopGpu::BeginSession]
ERROR [225876:233276] [[Profiling]] EnableMetrics returned -37 The third-party Profiling module returned an unrecoverable error [/workspaces/dcgm-rel_dcgm_2_3-postmerge/dc
gm_private/modules/profiling/DcgmModuleProfiling.cpp:2461] [DcgmNs::Modules::Profiling::DcgmModuleProfiling::ReconfigureLopGpu]
ERROR [225876:233276] [[Profiling]] Unable to reconfigure LOP metric watches for GpuId {0} [/workspaces/dcgm-rel_dcgm_2_3-postmerge/dcgm_private/modules/profiling/DcgmModu
leProfiling.cpp:2545] [DcgmNs::Modules::Profiling::DcgmModuleProfiling::ChangeWatchState]

导致这个错误的原因是对于这些 Profile 指标(1001-1014 ),NV 的 Profiler 对每个硬件使用了一个唯一锁。当你启动了超过 1 个的 nv-hostengine (包括内建的),比如使用 NCGM-Exporter 容器时,已经启动了一个内建的 nv-hostengine,然后又在主机上又启动了一个 nv-hostengine 服务,在访问这些指标时,就会出现这种访问失败。解决方案就是一台服务器仅启动一个 nv-hostengine 服务,然后所有的客户端都用本地或者远程的方式去访问它(5555 端口)。

这类问题在云场景下可能更常见,因为云服务商可能已经在租用的 GPU 服务器上部署了 DCGM 监控,你再去部署就可能遇到硬件锁的问题。

3. DCGM_FI_DEV_NVLINK_BANDWIDTH_TOTAL 的单位

我最早看到这个指标是在 NCGM-Exporter 的输出指标中,但是无论是文档、配置文件还是输出接口中,都没有写明这个指标的单位和计算逻辑。我问 Nvidia 的售后工程师,他告诉我这个指标的单位是 B/s,LoL!

后来仔细查 NCGM 的文档,发现 NCGM-Exporter 的所有指标其实都源自 NCGM 的接口,但还是没看到这个指标的单位和计算逻辑。

最后使用 dcgmi dmon -e 449,1011,1012 做了一下对比,才发现其输出头中有个不起眼的 MB 字样:

#Entity   NBWLT                       NVLTX                       NVLRX
ID       MB/

我做了一下数字的校验,基本可以明确 DCGM_FI_DEV_NVLINK_BANDWIDTH_TOTAL[449] 约等于 (DCGM_FI_PROF_NVLINK_TX_BYTES[1011] + DCGM_FI_PROF_NVLINK_RX_BYTES[1012]) / 1048576,所以单位应该是 MB/s 。

4. NCGM-Exporter 的采集延迟和精度

用 docker 跑 NCGM-Exporter 镜像时,发现指标的输出非常慢,输出的值看起来也很奇怪。后来自己构建了一下,研究了一下命令行参数,才发现默认采集的周期是 30s 一次。这种采样精度下,输出的指标值能准确才怪了。正常 dcgmi 采样的频率是 1s 一次,最低可以配置到 100ms 一次。NCGM-Exporter 有命令行参数可以调整这个采样频率,但需要你自己用修改 Dockerfile 去重建镜像。或者可以考虑启动的时候将启动脚本映射到外部文件?没尝试。

GPT-2 Tokenizer 效率观察

对基于 Transformer 结构的 LLM (大语言模型)来说,模型的输入输出都是 Token(词元)。一段输入文本,首先要经过 Tokenizer(分词器)切分成 Token 再输入给模型。

不同的 Tokenizer 会把文本按不同的边界切分,那一段文本到底会被切成几个 Token 就体现了 Tokenizer 本身的效率,这本身也是信息论的讨论范畴。不过今天不做理论分析,也不介绍 Tokenizer 的训练算法,就看下 GPT-2 和 GPT-3 Tokenizer 的实际分词效率(GPT-2 和 GPT-3 使用的是同样的 Tokenizer)。

用 GPT-2 Tokenizer 编码一段中文

from transformers import GPT2Tokenizer

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
input_str = '不同的 Tokenizer 会把文本按不同的边界切分,那一段文本到底会被切成几个 Token 就体现了 Tokenizer 本身的效率,这本身也是信息论的讨论范畴。'
tokenids = tokenizer(input_str)['input_ids']
raw_tokens = tokenizer.convert_ids_to_tokens(tokenids)
token_strs = [tokenizer.convert_tokens_to_string([token]) for token in raw_tokens]
print(len(input_str), len(tokenids), tokenids, token_strs)
82 121 [38834, 28938, 234, 21410, 29130, 7509, 220, 27670, 248, 162, 232, 232, 23877, 229, 17312, 105, 162, 234, 231, 38834, 28938, 234, 21410, 164, 122, 117, 45911, 234, 26344, 229, 26344, 228, 171, 120, 234, 165, 224, 96, 31660, 162, 106, 113, 23877, 229, 17312, 105, 26344, 108, 41753, 243, 27670, 248, 164, 95, 104, 26344, 229, 22755, 238, 49035, 254, 10310, 103, 29130, 10263, 108, 109, 19526, 241, 163, 236, 108, 12859, 228, 29130, 7509, 42164, 105, 164, 118, 104, 21410, 46763, 230, 163, 236, 229, 171, 120, 234, 32573, 247, 17312, 105, 164, 118, 104, 20046, 253, 42468, 46479, 94, 162, 223, 107, 164, 106, 118, 21410, 164, 106, 101, 164, 106, 118, 164, 234, 225, 45911, 112, 16764] ['不', '�', '�', '的', ' Token', 'izer', ' ', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '不', '�', '�', '的', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '一', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', ' Token', ' �', '�', '�', '�', '�', '�', '�', '�', '�', '�', ' Token', 'izer', ' �', '�', '�', '�', '�', '的', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '是', '�', '�', '�', '�', '�', '�', '�', '�', '的', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '�', '。']

输入的中英文夹杂的字符串长度为 82 个字符,其中有 6 个空格,23 个英文字符(3个单词),3个标点符号,50 个汉字字符。最终编码出来的 tokenids 长度是 121 个 Token,仔细检查可以看到,大部分汉字被编码成两个 Token,英文则是按照 subword 级别进行了编码。由此可见,GPT-2 Tokenizer 编码中文的效率是很低的

GPT-2 Tokenizer 词典中包含中文的 Token 数量

import re
tokens = [tokenizer.convert_tokens_to_string([key]) for key in tokenizer.get_vocab()]
cn_tokens = [t for t in tokens if re.search(u'[\u4e00-\u9fff]', t)]
cn_tokens.sort(key=len, reverse=True)
print(len(tokens), len(cn_tokens), cn_tokens)
50257 51 [' 裏覚醒', ' 裏�', '��士', '龍喚士', ' 裏�', '龍契士', '��極', ' 裏', '龍�', '�醒', '覚醒', 'の魔', '�士', '龍�', ' 神', '龍', '神', '士', '魔', '的', '人', '天', '王', '一', '大', '装', '田', '子', '戦', '生', '者', '不', '姫', '中', '上', '闘', '是', '女', '方', '作', '黒', '之', '使', '光', '代', '版', '三', '五', '武', '将', '機']

对 GPT-2 Tokenizer 词典中的 Token 进行分析,可以看到:在数量为 50257 大小的词典中,只有 51 个 Token 包含常见的中文字符。上一节中出现的“是”、“的”、“不”也在其中,但是大家随便一想就能想到的,常见的“你”、“我”、“他”并不在内。这些数据印证了它对中文分词的效率低下。

GPT-2 Tokenizer 词典中包含英文的 Token 数量

en_tokens = [t for t in tokens if re.search(u'[A-Za-z]', t)]
en_tokens.sort(key=len, reverse=True)
print(len(en_tokens), en_tokens[:10])
46949 ['rawdownloadcloneembedreportprint', 'BuyableInstoreAndOnline', 'cloneembedreportprint', ' RandomRedditorWithNo', ' telecommunications', ' disproportionately', ' guiActiveUnfocused', 'channelAvailability', ' Telecommunications', ' externalToEVAOnly']

但是针对英文字符进行统计,可以看到在数量为 50257 大小的词典中,有 46949 个 Token 包含 26 个英文字母。由于数量太多,这里仅输出了长度排前 10 的英文 Token。

从 Tokenizer 词典中不同语言字符的分布,我们可以验证 GPT-3 论文 “Language Models are Few-Shot Learners” 中提到的这句话:

... This could be a weakness due to reusing the byte-level BPE tokenizer of GPT-2 which was developed for an almost entirely English training dataset.

GPT-2 Tokenizer 词典中最长的 Token 是什么

tokens.sort(key=len, reverse=True)
print(len(tokens), tokens[:10])
50257 [' =================================================================', ' ----------------------------------------------------------------', '----------------------------------------------------------------', '................................................................', '================================================================', '________________________________________________________________', 'ÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂ', '--------------------------------------------------------', '------------------------------------------------', ' =================================']

可以看到,最长的 Token 居然是一些连续的标记,看起来像是论坛上的分割线。在“Natural Language Processing with Transformers” 一书中说,这可能是由于 GPT-2 的词典是用以 Reddit 为核心的语料训练的。

Codex Tokenizer 提升代码处理效率的 Trick

上面讲到 Tokenizer 效率对最终生成 Token 序列大小的影响,如果只关注最终的模型效果,可能并不在意中间过程的效率问题。

但从一个小点可以看到 OpenAI 还是很重视 Tokenizer 效率的。在训练 Codex 时,它在 50257 的词典之外新增了 23 个 Token,分别编码了:[2个连续空格, 3个连续空格, ..., 23个连续空格]。

大家都知道,代码中有很多锁进和对齐,如果每个空格都编码成一个 Token,训练和推理的效率就会极低;如果像常规语言模型那样忽视更多连续的空格,又会导致生成的代码不可读,甚至对 Python 这样的语言来说代码完全不可理解。所以 OpenAI 专门修改了 Tokenizer 的词典,提升了对连续空格的编码效率。

在中文数据集上训练的 Tokenizer 表现如何?

同样使用第一节的代码,我们换一个在中文数据集上训练的 Tokenizer,效果会如何呢?我随便在 HuggingFace 上找了一个模型 HuiHuang/gpt3-base-zh,使用它的 Tokenizer 对那段话进行重新分词。

from transformers import BertTokenizerFast

tokenizer = BertTokenizerFast.from_pretrained("HuiHuang/gpt3-base-zh")
input_str = '不同的 Tokenizer 会把文本按不同的边界切分,那一段文本到底会被切成几个 Token 就体现了 Tokenizer 本身的效率,这本身也是信息论的讨论范畴。'
tokenids = tokenizer(input_str)['input_ids']
raw_tokens = tokenizer.convert_ids_to_tokens(tokenids)
token_strs = [tokenizer.convert_tokens_to_string([token]) for token in raw_tokens]
print(len(input_str), len(tokenids), tokenids, token_strs)
82 65 [2, 215, 1750, 10574, 22090, 24330, 23635, 21748, 484, 5460, 6225, 6646, 5587, 215, 1750, 10574, 17027, 10262, 1233, 1232, 21124, 17261, 202, 7807, 6225, 6646, 1274, 4447, 484, 15221, 1233, 5338, 1194, 244, 22090, 24330, 3835, 541, 9850, 336, 22090, 24330, 23635, 21748, 6646, 16757, 10574, 6162, 9809, 21124, 17059, 6646, 16757, 297, 6393, 683, 4921, 16004, 10574, 15986, 16004, 13773, 10302, 160, 3] ['[CLS]', '不', '同', '的', 'to', '##ken', '##ize', '##r', '会', '把', '文', '本', '按', '不', '同', '的', '边', '界', '切', '分', ',', '那', '一', '段', '文', '本', '到', '底', '会', '被', '切', '成', '几', '个', 'to', '##ken', '就', '体', '现', '了', 'to', '##ken', '##ize', '##r', '本', '身', '的', '效', '率', ',', '这', '本', '身', '也', '是', '信', '息', '论', '的', '讨', '论', '范', '畴', '。', '[SEP]']

可以看到,82 个字符的输入被分成了 63 个 Token(除去[CLS]和[SEP]),比 GPT-2 Tokenizer 的分词数量 121 几乎少了一半。这就是 Tokenizer 训练数据集对分词效率的影响。

Tokenizer 的效率意味着什么

在以前,Tokenizer 的效率只是算法工程师们关注的话题。但是现在,LLM 逐渐成为云服务,Token 成为 API 调用的基本收费单元。例如:OpenAI gpt-3.5-turbo 的收费标准是 $0.002/1K tokens。

在这种情况下,Tokenizer 的效率就意味着成本。这意味着使用中文调用 OpenAI 的 chatGPT API,假设单 Token 成本一样的话,每次调用的实际成本可能是调用国内类似 chatGPT 云服务的 2 倍以上。

GPT-4 的 Bits Per Word 指标是什么意思

最近人人都在聊 ChatGPT,这里聊点儿更底层的技术细节。上周 OpenAI 发布了 GPT-4,里面有一张图引起了我的兴趣。在 OpenAI 内部的 Codebase 上,他们准确预测了 next word prediction 误差随着计算规模增加的下降曲线。让我感兴趣的不是这个曲线,而是这个曲线纵轴使用的指标:Bits Per Word (BPW)

图片来自:https://openai.com/research/gpt-4

乍一看我以为这是信息论里的概念,Bits Per Word 可能是指在某种语言里表达一个词所需要的最小位数,比如用于压缩时,这意味着压缩比的理论上限。但仔细看下指标接近 1.2 BPW,虽说英文字母的 Bits Per Character (BPC) 大约也在 1.2 左右,我不太相信单词的 BPW 能降到 1.2 这个规模。

后来才了解到,Bits Per Character (BPC) / Bits Per Word (BPW) 是 NLP 任务中对语言模型的评估指标之一。这些评估指标的详细解释可以先阅读这篇文章《Evaluation Metrics for Language Modeling》,但是我感觉这篇文章的说明在理论和应用之间跳过了一些环节,下面我以我的理解补全一下这些环节。推测的内容以黑字标出,如有谬误,欢迎指正。

原始 BPW 定义

The entropy is a statistical parameter which measures, in a certain sense, how much information is produced on the average for each letter of a text in the language. If the language is translated into binary digits (0 or 1) in the most efficient way, the entropy \(H\) is the average number of binary digits required per letter of the original language.

以上,是香农在《Prediction and Entropy of Printed English》一文中对语言的熵的定义。可以看到,语言的熵 \(H\) 其实就是 BPC (Bits Per Character) 的均值。但我们一般不会报告某个字符有几个 bits,因而通常所说的 BPC,就是整个语言(数据集)的平均 BPC,所以可以理解为语言的熵 \(H\) 就是 BPC。

但是在原始论文中,香农是用字符来建模的语言的熵,当用单词来建模语言时,语言的熵 \(H\) 就是 BPW (Bits Per Word)。

用 BPW 表示交叉熵

对于生成式语言模型来说,它的目标就是从样本中学习到一个分布 \(Q\),使这个分布尽量接近语言的经验分布 \(P\)。交叉熵 \(H(P, Q)\) 经常被用来衡量这种接近程度,

\[H(P, Q) = H(P) + D_{KL}(P \Vert Q) \]

根据 《Evaluation》 一文,交叉熵的两个因子用信息论的角度来理解:

  • \(H(P)\): 就是 \(P\) 的熵,即使用为分布 \(P\) 优化的编码器编码 \(P\) 中所有可能出现的字符的平均位数;
  • \(D_{KL}(P \Vert Q)\): 使用为分布 \(Q\) 优化的编码器编码 \(P\) 中所有可能出现的字符,所需要的平均额外位数;

与原始 BPC/BPW 的定义相同,如果用 \(BPC/BPW(P, Q)\) 表示交叉熵 \(H(P, Q)\),那么它意思应该是使用为分布 \(Q\) 优化的编码器编码 \(P\) 中所有可能出现的字符/单词所需要的平均位数。

上面这点理解,与知乎上一篇文章《一文搞懂Language Modeling三大评估标准》存在 diff(飘红部分):

也就是说BPC/BPW是cross-entropy对句子长度的平均,我们可以很容易地得出它的信息论含义:

基于预测的Q序列来编码P序列所需要的额外bit数在句子长度上的平均,也就是平均每一个字母/单词需要额外的bit数。

下面这一节,可能能解释 diff 的由来。

用 BPW 表示损失指标

但是在评估一个语言模型时,我们并不是评估模型输出的整个分布,而是评估模型的输出跟实际样本的不同。拿 GPT-2 模型来说,它是拿模型输出的预测下一个 Token 的 logits 与 label (样本中真实的下一个 Token)计算交叉熵:

The output from the decoder is passed to a language modeling head, which performs a linear transformation to convert the hidden states into logits. The label is the next token in the sequence, which are created by shifting the logits to the right by one. The cross-entropy loss is calculated between the shifted logits and the labels to output the next most likely token.

  • logits 是一个 \(Q(x)\) 的向量,表示下一个 Token 是 \(x\) 的概率;
  • label 表示样本中下一个 Token 是什么,它是一个 \(P(x)\) 的向量,但只有 x = 下一个 Token 时为 1,其它位置为 0,那么:

\[H(P) = -\sum_{x}P(x)\log{P(x)} = 0\]

从信息论也好理解,在计算损失时样本 \(P\) 是一个确定性分布,确定性分布的熵应该为 0。那么在计算交叉熵损失的场景下:

\[H(P, Q) = H(P) + D_{KL}(P \Vert Q) = 0 + D_{KL}(P \Vert Q)\]

可以看到,这时候交叉熵等于相关熵,那么 \(BPC/BPW(P, Q)\) 也就是 \(D_{KL}(P \Vert Q)\),或者更精确地说,是在数据集样本上平均的 \(D_{KL}(P \Vert Q)\)。那么在评估一个语言模型的场景下,把 BPC/BPW 损失指标的信息学含义解释成:“使用为分布 \(Q\) 优化的编码器编码 \(P\) 中所有可能出现的字符/单词所需要的平均额外位数”,也是合理的了。

GPT-4 的 BPW 指标

从 GPT-4 给出的指标来看,它的 BPW 指标只有 1.26 左右。除了信息论编码的概念,这代表什么意思呢?

这里我引入《Evaluation》 一文中提到的另一个指标 Perplexity (PPL) 困惑度和一个有趣的比喻。考虑到:

\[PPL(P, Q) = 2^{H(P, Q)} = 2^{BPW(P, Q)} = 2^{1.26} \approx 2.4\]

那么 GPT-4 语言模型的效果相当于:在给定的数据集上,提供一个样本的前缀串后,GPT-4 会制作一个包含下一个 Token 的 2.4 面的新骰子,然后掷骰子来预测下一个 Token 是什么。

趣谈哈希表优化:从规避 Hash 冲突到利用 Hash 冲突

本文首发于“百度架构师”公众号:https://mp.weixin.qq.com/s/oeuExiW3DYQnBG8HvDgBAg

1 背景

哈希表是一种查找性能非常优异的数据结构,它在计算机系统中存在着广泛的应用。尽管哈希表理论上的查找时间复杂度是 O(1),但不同的哈希表在实现上仍然存在巨大的性能差异,因而工程师们对更优秀哈希数据结构的探索也从未停止。

1.1 哈希表设计的核心

从计算机理论上来说,哈希表就是一个可以通过哈希函数将 Key 映射到 Value 存储位置的数据结构。那么哈希表设计的核心就是两点:

  1. 怎样提升将 Key 映射到 Value 存储位置的效率?
  2. 怎样降低存储数据结构的空间开销?

由于存储空间开销也是设计时的一个核心控制点,在受限于有限的空间情况下,哈希函数的映射算法就存在着非常高的概率将不同的 Key 映射到同一个存储位置,也就是哈希冲突。大部分哈希表设计的区别,就在于它如何处理哈希冲突。

当遇到哈希冲突时,有几种常见的解决方案:开放寻址法、拉链法、二次哈希法。但是下面我们介绍两种有趣的、不常见的解决思路,并且引出一个我们新的实现方案 - B16 哈希表

2 规避哈希冲突

传统哈希表对哈希冲突的处理会增加额外的分支跳转和内存访问,这会让流水线式的 CPU 指令处理效率变差。那么肯定就有人考虑,怎么能完全规避哈希冲突?所以就有了这样一种函数,那就是完美哈希函数(perfect hash function)。

完美哈希函数可以将一个 Key 集合无冲突地映射到一个整数集合中。如果这个目标整数集合的大小与输入集合相同,那么它可以被称为最小完美哈希函数。

完美哈希函数的设计往往非常精巧。例如 CMPH (http://cmph.sourceforge.net/)函数库提供的 CDZ 完美哈希函数,利用了数学上的无环随机 3-部超图概念。CDZ 通过 3 个不同的 Hash 函数将每个 Key 随机映射到 3-部超图的一个超边,如果该超图通过无环检测,再将每个 Key 映射到超图的一个顶点上,然后通过一个精心设计的与超图顶点数相同的辅助数组取得 Key 最终对应的存储下标。

完美哈希函数听起来很优雅,但事实上也有着实用性上的一些缺陷:

  • 完美哈希函数往往仅能作用在限定集合上,即所有可能的 Key 都属于一个超集,它无法处理没见过的 Key;
  • 完美哈希函数的构造有一定的复杂度,而且存在失败的概率;
  • 完美哈希函数与密码学上的哈希函数不同,它往往不是一个简单的数学函数,而是数据结构+算法组成的一个功能函数,它也有存储空间开销、访问开销和额外的分支跳转开销;

但是在指定的场景下,例如只读的场景、集合确定的场景(例如:汉字集合),完美哈希函数可能会取得非常好的表现。

3 利用哈希冲突

即便不使用完美哈希函数,很多哈希表也会刻意控制哈希冲突的概率。最简单的办法是通过控制 Load Factor 控制哈希表的空间开销,使哈希表的桶数组保留足够的空洞以容纳新增的 Key。Load Factor 像是控制哈希表效率的一个超参数,一般来说,Load Factor 越小,空间浪费越大,哈希表性能也越好。

但近年来一些新技术的出现让我们看到了解决哈希冲突的另一种可能,那就是充分利用哈希冲突

3.1 SIMD 指令

SIMD 是单指令多数据流(Single Instruction Multiple Data)的缩写。这类指令可以使用一条指令操作多个数据,例如这些年非常火的 GPU,就是通过超大规模的 SIMD 计算引擎实现对神经网络计算的加速。

目前的主流 CPU 处理器已经有了丰富的 SIMD 指令集支持。例如大家可接触到的 X86 CPU 大部分都已经支持了 SSE4.2 和 AVX 指令集,ARM CPU 也有 Neon 指令集。但科学计算以外的大部分应用程序,对 SIMD 指令的应用还都不太充分。

3.2 F14 哈希表

Facebook 在 Folly 库中开源的 F14 哈希表有个非常精巧的设计,那就是将 Key 映射到块,然后在块里使用 SIMD 指令进行高效过滤。因为分块的数量比传统的分桶要更小,这相当于人为增加了哈希冲突,然后在块中用 SIMD 指令再解决冲突。具体的做法是这样的:

  • 通过哈希函数对 Key 计算出两个哈希码:H1H2, 其中 H1 用来确定 Key 映射到的块,H2 只有 8 bits,用来在块内进行过滤;
  • 每个块里最多存放 14 个元素,块头有 16 个字节。块头的前 14 个字节,存放的是 14 个元素对应的 H2,第 15 字节是控制字节,主要记录该块里有多少个元素是从上一个块溢出的,第 16 字节是越界计数器,主要记录如果该块空间足够大,应该会被放置多少个元素。
  • 在插入时,当 Key 映射到的块中 14 个位置还有空时,就直接插入;当块已经满时,就增加越界计数器,尝试插入下一个块中;
  • 在查询时,通过待查找 Key 计算得到 H1H2。通过 H1 对块数取模确定其所属的块后,首先读取块头,通过 SIMD 指令并行比较 H2 与 14 个元素的 H2s 是否相同。如果有相同的 H2,那么再比对 Key 是否相同以确定最终结果;否则根据块头的第 16 字节判断是否还需要比对下一个块。

F14 为了充分利用 SIMD 指令的并行度,在块内使用了 H2 这种 8 bits 的哈希值。因为一个 128 bits 宽度的 SIMD 指令可以进行最多 16 个 8 bits 整数的并行比较。虽然 8 bits 哈希值的理论冲突概率 1/256 并不低,但也相当于有 255/256 的可能性省去了逐个 Key 对比的开销,使哈希表能够容忍更高的冲突概率。

4 B16 哈希表

不考虑分块内部的设计,F14 本质上是一个开放寻址的哈希表。每个块头的第 15、16 字节被用于开放寻址的控制策略,只剩 14 个字节留给哈希码,也因而被命名为 F14。

那么我们考虑能不能从另一个角度出发,使用拉链法来组织分块。由于省去了控制信息,每个分块中可以放置 16 个元素,我们把它命名为 B16。

4.1 B16 哈希数据结构

b16

图 1:B16 哈希表数据结构(3元素示例)​

上图所示就是用每个分块 3 个元素展示的 B16 哈希表的数据结构。中间绿色的是常见的 BUCKET 数组,存放的是每个分桶中 CHUNK 拉链的头指针。右侧的每个 CHUNK 与 F14 相比,少了控制字节,多了指向下一个 CHUNK 的 next 指针。

B16 也是通过哈希函数对 Key 计算出两个哈希码:H1H2。例如 “Lemon” 的两个哈希码是 0x24EB 和 0x24,使用 H1 的高位作为 H2 一般来说足够了。

在插入时,通过 H1 对桶大小取模计算 Key 所在的分桶,例如 "Lemon" 所在的分桶是 0x24EB mod 3 = 1。然后在 1 号分桶的分块拉链中找到第一个空位,将 Key 对应的 H2 和元素写入该分块。当分块拉链不存在,或者已经装满时,为拉链创建一个新的分块用于装载插入的元素。

在查找时,先通过 H1 找到对应的分桶拉链,然后对每个块进行基于 SIMD 指令的 H2 对比。将每个块的块头 16 字节加载到 128 bits 寄存器中,里面包含了 16 个 H2',把 H2 也重复展开到 128 bits 寄存器中,通过 SIMD 指令进行 16 路同时比对。如果都不同,那就对比下一个块;如果存在相同的 H2,就继续对比对应元素的 Key 是否与查找的 Key 相同。直到遍历完整条拉链,或者找到对应的元素。

在删除时,首先查找到对应的元素,然后用块拉链最尾部的元素覆盖掉对应的元素即可。

当然,B16 哈希表每个块的元素个数可以根据 SIMD 指令的宽度进行灵活调整,例如使用 256 bits 宽度指令可以选择 32 元素的块大小。但影响哈希表性能的不仅仅是查找算法,内存访问的速度和连续性也非常重要。控制块大小在 16 以内在大多数情况下能充分利用 x86 CPU 的 cache line,是一个较优的选择。

普通的拉链式哈希表,拉链的每个节点都只有一个元素。B16 这种分块式拉链法,每个节点包含 16 个元素,这会造成很多空洞。为了使空洞尽可能的少,我们就必须增加哈希冲突的概率,也就是尽量地缩小 BUCKET 数组的大小。我们经过试验发现,当 Load Factor 在 11-13 之间时,B16 的整体性能表现最好。其实这也相当于把原来存在于 BUCKET 数组中的空洞,转移到了 CHUNK 拉链中,还省去了普通拉链式每个节点的 next 指针开销。

4.2 B16Compact 哈希数据结构

为了进一步压榨哈希表的存储空间,我们还设计了 B16 的只读紧凑数据结构,如下图所示:

b16

图 2:B16Compact 哈希表数据结构(3元素示例)​

B16Compact 对哈希表结构做了极致的压缩。

首先它省去了 CHUNK 中的 next 指针,把所有的 CHUNK 合并到了一个数组中,并且补上了所有的 CHUNK 空洞。例如【图 1】中 BUCKET[1] 的拉链中本来有 4 个元素,包含 Banana 和 Lemon,其中头两个元素被补到了【图 2】的 CHUNK[0] 中。以此类推,除 CHUNK 数组中最后一个 CHUNK 外,所有的 CHUNK 均是满的。

然后它省去了 BUCKET 中指向 CHUNK 拉链的指针,只保留了一个指向原拉链中第一个元素所在 CHUNK 的数组下标。例如【图 1】中 BUCKET[1] 的拉链第一个元素被补到了【图 2】的BUCKET[0] 中,那么新的 BUCKET[1] 中仅存储了 0 这个下标。

最后增加了一个 tail BUCKET,记录 CHUNK 数组中最后一个 CHUNK 的下标。

经过这样的处理以后,原来每个 BUCKET 拉链中的元素在新的数据结构中依然是连续的,每个 BUCKET 依然指向第一个包含其元素的 CHUNK 块,通过下一个 BUCKET 中的下标依然可以知道最后一个包含其元素的 CHUNK 块。不同的是,每个 CHUNK 块中可能会包含多个 BUCKET 拉链的元素。虽然可能要查找的 CHUNK 变多了,但由于每个 CHUNK 都可以通过 SIMD 指令进行快速筛选,对整体查找性能的影响相对较小。

这个只读哈希表只支持查找,查找的过程跟原来差异不大。以 Lemon 为例,首先通过 H1=24EB 找到对应的分桶 1,获得分桶对应拉链的起始 CHUNK 下标为 0,结束 CHUNK 下标为 1。使用与 B16 同样的算法在 CHUNK[0] 中查找,未找到 Lemon,然后继续查找 CHUNK[1],找到对应的元素。

B16Compact 的理论额外存储开销可以通过下式算得:

其中,n 是只读哈希表的元素个数。

当 n 取 100 万,Load Factor 取 13 时,B16Compact 哈希表的理论额外存储开销是 9.23 bits/key,即存储每个 key 所支出的额外开销只有 1 个字节多一点。这几乎可以媲美一些最小完美哈希函数了,而且不会出现构建失败的情况。

B16Compact 数据结构仅包含两个数组,BUCKET 数组和 CHUNK 数组,也就意味着它的序列化和反序列化可以做到极简。由于 BUCKET 中存储的是数组下标,用户甚至不需要将哈希表整个加载到内存中,使用文件偏移即可进行基于外存的哈希查找,对于巨大的哈希表来说可以有效节省内存。

5 实验

5.1 实验设定

实验使用的哈希表的 Key 和 Value 类型均取 uint64_t,Key、Value 对的输入数组由随机数生成器预先生成。哈希表均使用元素个数进行初始化,即插入过程中不需要再 rehash。

  • 插入性能:通过 N 个元素的逐一插入总耗时除以 N 获得,单位是 ns/key;
  • 查询性能:通过对随机的 Key 查询 20 万次(全命中) + 对随机的 Value 查询 20 万次(有可能不命中)获得总耗时除以 40 万获得,单位是 ns/key;
  • 存储空间:通过总分配空间除以哈希表大小获得,单位是 bytes/key。对于总分配空间来说,F14 和 B16 均有对应的接口函数可直接获取,unordered_map 通过以下公式获取:
// 拉链节点总大小
umap.size() * (sizeof(uint64_t) + sizeof(std::pair<uint64_t, uint64_t>) + sizeof (void*)) 
// bucket 总大小
+ umap.bucket_count() * sizeof (void *)
// 控制元素大小
+ 3 * sizeof(void*) + sizeof(size_t);

Folly 库使用 -mavx -O2 编译,Load Factor 使用默认参数;B16 使用 -mavx -O2 编译,Load Factor 设定为 13;unordered_map 使用 Ubuntu 系统自带版本,Load Factor 使用默认参数。

测试服务器是一台 4 核 8G 的 CentOS 7u5 虚拟机,CPU 是 Intel(R) Xeon(R) Gold 6148 @ 2.40GHz,程序在 Ubuntu 20.04.1 LTS Docker 中编译执行。

5.2 实验数据

b16

图 3:插入性能对比​

上图中折线所示为 unordered_map、F14ValueMap 和 B16ValueMap 的插入性能对比,不同的柱子显示了不同哈希表的存储开销。

可以看到 B16 哈希表在存储开销明显小于 unordered_map 的情况下,仍然提供了显著优于 unordered_map 的插入性能。

由于 F14 哈希表对 Load Factor 的自动寻优策略不同,在不同哈希表大小下 F14 的存储空间开销存在一定波动,但 B16 的存储开销整体仍优于 F14。B16 的插入性能在 100 万 Key 以下时优于 F14,但是在 1000 万 Key 时劣于 F14,可能是因为当数据量较大时 B16 拉链式内存访问的局部性比 F14 差一些。

b16

图 4:查找性能对比​

上图中折线所示为 unordered_map、F14ValueMap、B16ValueMap 和 B16Compact 的查找性能对比,不同的柱子显示了不同哈希表的存储开销。

可以看到 B16、B16Compact 哈希表在存储开销明显小于 unordered_map 的情况下,仍然提供了显著优于 unordered_map 的查找性能。

B16 与 F14 哈希表的查找性能对比与插入性能类似,在 100 万 Key 以下时明显优于 F14,但在 1000 万时略劣于 F14。

值得注意的是 B16Compact 哈希表的表现。由于实验哈希表的 Key 和 Value 类型均取 uint64_t,存储 Key, Value 对本身就需要消耗 16 字节的空间,而 B16Compact 哈希表对不同大小的哈希表以稳定的约 17.31 bytes/key 进行存储,这意味着哈希结构里为每个 Key 仅额外花费了 1.31 个字节。之所以没有达到 9.23 bits/key 的理论开销,是因为我们的 BUCKET 数组没有使用 bitpack 方式进行极致压缩(这可能会影响性能),而是使用了 uint32_t。

6 总结

本文中我们从哈希表设计的核心出发,介绍了两种有趣的、不算“常见”的哈希冲突解决方法:完美哈希函数和基于 SIMD 指令的 F14 哈希表。

在 F14 的启发下,我们设计了 B16 哈希表,使用了更容易理解的数据结构,使得增、删、查的实现逻辑更为简单。实验表明在一些场景下 B16 的存储开销和性能比 F14 还要好。

为追求极致的存储空间优化,我们对 B16 哈希表进行紧致压缩,设计了几乎可以媲美最小完美哈希函数的 B16Compact 哈希表。B16Compact 哈希表的存储开销显著低于 F14 哈希表(介于40%-60%之间),但却提供了颇有竞争力的查询性能。这在内存紧张的场合,例如嵌入式设备或者手机上,可以发挥很大的作用。

新的哈希表设计表明 SIMD 指令的并行化处理能力的有效应用能大幅度提升哈希表对哈希冲突的容忍能力,进而提升查询的速度,并且能帮助哈希表进行极致的存储空间压缩。这让哈希表的设计思路从尽量规避哈希冲突,转向了利用合适的哈希冲突概率来优化计算和存储效率。

使 Netron 支持 PaddlePaddle 模型子图显示

先介绍一下什么是 PaddlePaddle 的模型子图。一般的神经网络都可以表示成一张由算子组成的计算图,但是对一些较为复杂的神经网络,可能会存在一些条件分支。PaddlePaddle 在构建这种条件分支的网络时,会把分支内的计算图单独保存成一张子图。

具体到 PaddlePaddle 的内部数据结构时,每个子图就是 Program 内的一个 Block。这个 Block 内包含该子图的所有中间变量,op 算子和参数等。

目前 PaddlePaddle 的模型可视化主要是依赖 Netron,包括 PadddlePaddle 提供的 VisualDL 工具里内嵌的也是 Netron。

我最近在研究 PaddlePaddle 模型优化的时候,发现 VisualDL 完全不支持子图的可视化。我稍微研究了一下 Netron 的代码,想把子图的可视化给加上。但没想到的是,其实 Netron 本身已经做了 PaddlePaddle 子图结构的解析,只是代码上存在一点儿 bug,导致子图无法被选择出来。

我就修复了一下这个 BUG,提交了一个 PR: https://github.com/lutzroeder/netron/pull/588 。希望对 PaddlePaddle 的用户能有所帮助。

只是 Netron 项目的 Owner lutzroeder 有些奇怪。他不是直接接受这个 PR,而是自己做一些修改后提交一个新的 commit 引用了这个 PR,不知道出于什么目的。这个习惯会让人很难分析出来对 PaddlePaddle 模型解析代码感兴趣的人都是谁,对于维护并不便利。

基于 SIMD 指令的 PFOR-DELTA 解压和查找

PFOR-DELTA 是一种经典的有序整数数列压缩算法,被广泛使用在搜索、推荐引擎的倒排索引和召回队列压缩中。PFOR-DELTA 的具体算法这里就不展开了,不了解的同学可以参考它的原始论文《Super-Scalar RAM-CPU Cache Compression》或者做一些搜索工作。

朴素的 PFOR-DELTA 解压主要是逐个对 frame 中的 bitpack 整数进行解压,对朴素 PFOR-DELTA 的优化主要包括对齐的内存访问aligned memory access, 先按4/8字节偏移读取,然后再移位取得 bitpack)和循环展开loop unrolling, 由每次循环解压一个整数转成每次循环解压 N 个整数以减少分支判断和访存次数)。通用的 PFOR-DELTA 函数库,往往采取这两种优化方法。

本文主要介绍在目前的先进 CPU 架构下,如何利用 SIMD 指令(如 SSE, AVX) 加速 PFOR-DELTA 解压和查找。中文互联网上与之相关的有一篇阿里搜索和推荐团队的文章《索引压缩算法New PForDelta简介以及使用SIMD技术的优化》,但该文章缺乏算法细节且其收益表明对 SIMD 指令的应用并不高效。与该文类似,为简化计,下文主要以不带异常段的 PFOR-DELTA 为例来说明算法细节。

基于 SIMD 的 bit unpacking

PFOR-DELTA 的每个分块中,都是以固定位宽压缩的整数。例如小于 32 的整数,都可以用 5 位来存储,相比于原来的 32 位存储,大大减少了数据的存储空间。但是在使用的时候,我们又必须将压缩的数展开到 32 位,才能进计算和比较操作,将一个数从压缩的位宽展开到使用的位宽,叫做 bit unpacking

对单独一个 bitpack 的整数来说,展开是非常容易的,通过简单的移位、AND 操作即可完成。但是 SIMD 指令主要提升的是并行化,考虑到 frame 宽度不同,存在各种对齐问题,如何同时进行多个数的 bit unpacking 并不是一件非常直观的事。

下面以 frame 宽度为 9,即每个整数用 9 bit 表示,来详细说明基于现代 CPU 广泛支持的 128位 SIMD 指令的 bit unpacking 算法。256/512 位的 SIMD 算法可依此推演。

如上图所示,如果我们从数组开头加载 128 bit 的数据,其中将会包括 14 个完整的 9 bit 整数,以及多读的 2 个 bit。

为了将每个 9 bit 展开成 32 bit,在只有 128 位寄存器的情况下,我们只能 4 个 4 个地展开。首先需要将头 4 个 9 bit 整数移动到寄存器 32 位偏移的位置。Intel 提供了 _mm_shuffle_epi8 intrinsic 可以根据重整参数重新排布 128 位寄存器的内容,但可惜的是它最小粒度只能按 byte 重整,这也就意味着我们必须将 9 bit 整数所在的整个 2 byte 移动到 32 位偏移的位置。对应的代码是:

// 考虑到 litten endian,实际上 shufkey 内容为:
// 00018080 01028080 02038080 03048080
__m128i shufkey = {0x8080020180800100, 0x8080040380800302};
v = _mm_shuffle_epi8(v, shufkey);

这时候,128 位寄存器中的每个 32 位槽位中都包含一个 9 bit 整数,但可惜的是它都包含了前后数字的一些冗余 bit。下面这张图,说明了如何用 SIMD 指令消除这些冗余 bit。

前面说过,对一个整数做 unpack 可以完全用移位和 AND mask 来实现,对多个整数依然如此。但由于各槽位中 9 bit 整数的对齐不同,移位 bit 数也不同。例如:第一个槽位右移 0 位,第二个槽位右移 2 位,第三个操作右移 2 位,第四个槽位右移 3 位。

可是 SSE 的移位指令只能支持各槽位同等位数的移位,无法支持这样各槽位不同位数的移位。我们不得不用代价更高昂的其它指令来实现分槽位不同移位,先用基于 2 次幂的向量乘法实现左移对齐 9 bit 整数,然后再统一右移,统一 mask。不使用 2 次幂除法的原因是除法的成本更高。对应到代码是:

// 通过向量乘法实现各 DW 槽位中 9 bit 整数左移对齐。当然参数可以换成常数。
v = _mm_mullo_epi32(v, _mm_set_epi32(8, 4, 2, 1));
// 所有 DW 右移 3 位
v = _mm_srli_epi32(3);
// 通过 and mask 掉 9 bit 之外多余的 bit。当然参数可以换成常数。
v = _mm_and_si128(v, _mm_set1_epi32(0x1ff));

这样,我们就实现了 9 bit 整数的 unpacking。上述算法具备通用性,对于其它的 frame 宽度,只是 shuffle、移位、乘法、mask 的参数不同,处理过程并无区别。

认真的读者可能会有疑问:第一张图为什么做了两遍 shuffle?这是考虑到了奇数位 bit packing 的特殊性。

对于偶数位(2,4,6,8)的 bit packing,一次解压 4 个整数,解压的位宽是可以被 8 整除的,所以每次解压都可以从字节边界开始,所有算法参数都相同;对于奇数位(3, 5, 7, 9),一次解压 4 个整数,解压的位宽是不能被 8 整除的,所以第二轮解压不能从字节边界开始,第二轮解压的算法参数与第一轮不同。但解压 8 个整数后,位宽就对齐到字节边界了。所以说:frame 位宽为偶数的 SIMD unpacking,循环 block 可以是 4*位宽;但 frame 位宽为奇数的 SIMD unpacking,循环 block 只能是 8*位宽。如果 AVX 指令可用,可以使用 256 位的 SIMD 指令,能大大缓解这个问题。

不失其一般性,可以看到我们只用了 4 条 SIMD 计算指令,就完成了 4 个整数的 bit unpacking,其效率较逐一 unpacking 高了很多,并且可以推及到 256 位和 512 位的情形得到更高加速比。

本节内容主要参考了论文《SIMD-Scan: Ultra Fast in-Memory Table Scan using on-Chip Vector Processing Units》。

基于 SIMD 的 delta 计算

在 PFOR-DELTA 算法中,每个整数其实是 delta,需要把所有前序的整数加起来才是真正所要的数据。参考上节,我们虽然解压出来了 {v0, v1, v2, v3},但实际上需要的却是 {v0, v0+v1, v0+v1+v2, v0+v1+v2+v3}

将 __m128i 转成一个数组,然后再循环相加是比较简单的思路。但我们也可以直接用三条 SIMD 指令来实现 delta 计算:

// {v0, v1, v2, v3} + {0, 0, v0, v1} = {v0, v1, v0+v2, v1+v3}
v = _mm_add_epi32(_mm_slli_si128(v, 8), v);
// {v0, v1, v0+v2, v1+v3} + {0, v0, v1, v0+v2} = {v0, v0+v1, v0+v1+v2, v0+v1+v2+v3}
v = _mm_add_epi32(_mm_slli_si128(v, 4), v);
// {v0, v0+v1, v0+v1+v2, v0+v1+v2+v3} + {acc3, acc3, acc3, acc3}
acc = _mm_add_epi32(v, _mm_shuffle_epi32(acc, 0xff));

同样,不失一般性,这种错位相加的方法完全可以推广到 8 维整数向量的 delta 计算,可以推广到 256 位和 512 位的情形得到更高的加速比。

本节内容主要参考了论文《SIMD Compression and the Intersection of Sorted Integers》。

基于 SIMD 的查找比较

对 PFOR-DELTA 解压完成之后对有序数组,如果想找到某个整数的位置,我们还需要逐个进行比较。这种比较,也可以用 SIMD 指令完成:

// 初始化一次
__m128i key = _mm_set1_epi32(v_to_find);
// 向量比较
v = _mm_cmplt_epi32(key, acc);
// 比较结果处理
int res = _mm_movemask_epi8(v);
if (res != 0) {
index += __builtin_ctz(res) >> 2;
} else {
index += 4;
}

总结

综上,本文描述了完全使用 SIMD 指令进行 PFOR-DELTA 解压和查找的详细算法,给出了在 SSE 指令集下的具体代码,并且可以推广到更高的数据宽度下。至于优化的收益,将根据基线实现的不同存在差异,感兴趣的读者可以自行实现比较一下。

此外,上文的实现主要着眼通用性,针对特定的小宽度整数,其实可以使用更小的计算粒度以增大并行度。对性能有苛求的读者可自行研究。

ABTest 平台设计 - 流量分布问题

通过前三篇文章,大家可以了解到一个基本的 ABTest 平台架构建设的要点,看起来构建起一个生产环境可用的 ABTest 平台难度不大。但如果想把这个平台做得更强大,还有很多细节地方要注意。下面举两个常见的细节问题。

跨层分桶不完全正交

理论上来讲,如果每层分桶都是采取的独立不相关的 hash 算法,那么层和层之间的流量应该是完全正交且服从均匀分布的。举个例子:如果每层分 10 个桶,那么第二层的 1 号桶内,应该有第一层每个桶里的 1/10 的用户。

但从实际上来讲,很多情况下我们使用的 hash 算法没有那么地“强”,会导致分桶没那么地随机和均匀。比如几年前非常流行的 MurmurHash,也被 SipHash 的作者找到了攻击的方法,他们更推荐使用更强的 SipHash / HighwayHash

还有就是均匀分布仅仅是统计意义上的期望,而不是实际的效果。比如我们有 1-10 十个数,用 hash 算法一定能保证把它们分到 10 个桶里,每个桶里 1 个数吗?

只不过大多时候,在统计意义上能保证基本的均匀,对大部分产品来说也够用了。

对于一些强迫症公司,或者用户数还没有达到统计意义上可分又开始早早做 ABTest 的公司,或者从数据观察的确存在较大的实验间干涉的公司,可能会做这样的一个事,人工构造层间正交分桶

具体做法是这样的:先用 hash 算法(或其它抽样算法)将用户分成桶数的平方个最小流量单元,然后再选取这些流量单元组成每个分层。比如要构造 2 层,每层 3 个分桶的用户分层,先将用户分为 3*3 个流量单元,那么两层就可以这样构造:

Layer 1: [1, 2, 3] [4, 5, 6] [7, 8, 9]
Layer 2: [1, 4, 7] [2, 5, 8] [3, 6, 9]

可以看到第二层的每一个桶,都包含第一层里所有桶的 1/3 用户。通过精巧地构造来实现层间的正交分桶,可以有效地降低层间实验互相干涉的情况。

发版造成指标波动

当实验本身受到 APP 版本一定影响的时候,AB 分组可能会由于升级节奏不同,导致 AB 指标波动完全不可比。

举一个例子:如果新版本使用时长会上升 10%,正常情况下实验分组 A 比分组 B 使用时长上升 3%。那么如果分组 A 用户升级新版本占比 30%,分组 B 用户升级新版本占比 70%,结果会是怎样呢?

分组 A:70*1.03 + 30*1.1*1.03 = 106.09
分组 B:30 + 70*1.1 = 107

分组 A 反而比分组 B 的数据更差!当然,这里为了效果,举的例子比较极端。实际的情况可能是在发版收敛期间,ABTest 的指标波动较大。

而且升级动作本身对活跃用户和非活跃用户有一定的筛选作用,产生的行为和数据也是有偏的,很难有完美的解决办法。避开发版收敛时段,或者对数据进行多维的分析和组织,也许有一定帮助。

结语

按照我以前的风格,这本应是一篇文章。但在信息快消时代,我发现自己也没耐心读长文了,所以就拆成一个系列发布了。

这篇文章主要从平台建设的角度出发,讨论了一些平台设计中要考虑到的关键功能点,希望能对读者有所助益。至于如何科学地进行 ABTest 实验设计和效果分析,不是我擅长的部分,就不再后续展开了。

ABTest 平台设计 - 灰度发布和早鸟用户

上篇《ABTest 平台设计 - 实验开关和分组信息传递》简单介绍了 ABTest 实验开关和数据收集的一些实现,从流量划分、到实验开关、到数据收集,基本实现了 ABTest 的主要功能。下面则扩展谈一下 ABTest 的衍生功能:

基于 ABTest 的灰度发布

与 ABTest 一样,灰度发布也是圈出来一部分流量进行新功能的线上验证,验证基本能力没有问题之后再逐渐扩大覆盖面,支持扩展到全流量。

灰度发布本身也有很多种机制,例如最常用的:上线时先上单副本,再扩展到多副本,再扩展到单机房,再扩展到其它机房。这种方式非常自然,逐步扩量观察保证了服务稳定性。

但这样在上线的中间过程中,总不可避免地会出现一些用户体验问题。比如用户相邻的多次刷新请求被路由到版本不同的副本上,导致请求结果的跳变

既然 ABTest 同样具备划分流量的能力,而且这种划分对于单个用户来说是稳定的,其实在很多情况下可以利用 ABTest 能力来实现灰度发布。

但基于 ABTest 的灰度发布,要求在架构上提供一些支持

比如要发布新版本的网站静态文件(css/js 等),可以全量发布多个版本,然后通过 ABTest 圈定部分用户路由到新版本的静态文件,其余则路由到原版本的静态文件。

比如要发布新版本的服务程序,可以用新的 Docker 部署新版本程序,通过 ABTest 圈定部分用户请求路由到新的 Docker 上,其余则路由到线上的 Docker 上。

基于 ABTest 的灰度发布,很多情况下可以简化服务的部署和回滚操作,也保证了用户在灰度上线期间的体验稳定性。

服务好早鸟用户

在很多时候,企业的内外部总存在着一些早鸟用户,他们对灰度 / AB 新功能有着非常迫切的需求。

最典型的早鸟用户,就是公司的老板。当你开发了一个新功能,向上汇报了一下这功能多好多好。老板会问:

“为什么我没有看到这个新功能?”

你可能不得不解释说:

“老板你的 ID 没有被随机分到实验组里。”,

或者:

“老板这个功能只上了广州机房,北京访问不了。”

还有一种早鸟用户,是产品经理、测试人员,甚至可能包括提交 BUG 的外部用户。他们需要去回归新功能的线上效果是否达到预期,而且甚至他们需要一直不停地在不同的 AB 分支上切来切去比较效果。

这时候就需要一种灵活的机制,让这些早鸟用户有办法切换 AB 功能。

很多人最直接想到的方法,是在随机分桶之外搞一个 ID 分桶,收集上来早鸟用户的 ID,手动配置实验分组。这能在一定程度上解决问题,但设想这样的场景:

“老板,把你 UserID 发来一下,我给你配一下小流量实验。”

老板心里肯定在嘀咕,我要给了你 UserID,岂不是我看过啥发过啥你都能查出来了?这以后还能有隐私么?

而且配置 ID 的方式会增加运维的工作量,尤其是用户需要切来切去的时候。所以这时候不如提供一些强制命中灰度 / AB 的后门功能,能够大幅度降低沟通和维护的复杂度。

这种给早鸟用户的后门可以有很多种做法:

写 Cookie 机制。提供一个特殊的 URL,访问该 URL 就会种下一个强制命中的实验分组 Cookie,此后带着这个分组 Cookie 的访问都会中这个实验。这种适用于 WEB 端产品。APP 端使用的话,需要做一些 Cookie 同步工作。

配置注入机制。提供一个二维码,二维码内容是一段特殊的代码,APP 扫描到该二维码,就会被注入实验分组配置。这种适用于 APP 端产品。

隐藏功能机制。在某些内容上连续点 N 下,就会弹出一个配置面板,可以用来查看和调整当前所中的实验分组。

ABTest 平台如果能够提供这样的后门机制,将会大大方便与早鸟用户的沟通和合作。

下篇,我们聊一下流量分布问题。

ABTest 平台设计 - 实验开关和分组信息传递

通过上文《ABTest 平台设计 - 如何进行流量分桶》可以知道如何把用户流量科学地分配到不同的实验分组中,下面就面临一个问题:如何根据分组信息控制产品功能?

一种直观做法

最显而易见的做法,是直接在系统中传递分组信息,同时使用分组信息作为实验功能的开关

比如说在服务端系统的接入层,通常是 Nginx 或者其它入口模块,对流量进行分组,为每个请求添加一个抽样分组字段,字段内容类似于 “Exp1Group1,Exp3Group2,...”。然后所有对下游的请求,都带着这个分组字段。那么下游的各个模块,都可以根据这个分组字段来决定程序逻辑。用伪码表示就是:

if 抽样分组字段 包含 Exp1Group1:
do 实验1 的 A逻辑
elif 抽样分组字段 包含 Exp1Group2:
do 实验1 的 B逻辑
else
do 全流量逻辑

上面的伪码有些缺点,就是抽样分组信息 Exp*Group* 写死在了代码里,这样灵活性太差。所以,即使直接使用分组字段,通常也是将分组字段作为配置项写到代码里,这样更方便测试和部署。

客户端带来的挑战

在服务端直接使用抽样分组信息作为实验功能开关尚且可以忍受,因为一旦有问题调整下重新上线并不困难,但是在客户端 APP 中这样做就行不通了。客户端版本一旦发布,再想改动就面临着诸多的问题,比如发布审核周期问题,用户拒绝升级问题等等。

还有一个问题就是冗余代码和数据控制的问题。服务端功能 ABTest 转全以后,可以同时删除实验分组和对应的实验功能代码。客户端的实验代码是很难删除的,而且无法预判哪个分支会转全,所以没法直接删除实验分组。这会导致线上的实验越来越多,无法控制。

功能配置和分组配置分离的设计

回过头来思考这个需求:用户中了实验分组 A ,就走实验 A 逻辑。 其实可以加一层抽象:用户中了实验分组 A,就下发 A 对应的功能配置,实现走实验 A 逻辑

就拿客户端 ABTest 来说,大实验可能是换一下 APP 版式,增加或者减少一个新功能;小实验可能是调整一下字体和字号,调整一下背景颜色或者图片等。其实本质上是控制一些功能的配置项。ABTest 可能会下线,但这些配置项会一直在产品中存在。

功能配置和分组配置分离还能带来很大的灵活性,也就是说云端可以创建新的 ABTest 尝试不同的功能配置组合,而不需要硬编码固定的 ABTest。

假设你本来有两个实验,标题大小实验和内容大图实验。如果硬编码情况下,你只能做原始实验方案的对比。如果用功能配置的话,等这两个实验完了,你还可以实验不同标题和内容大小图的结合实验,这是不需要再开发和发版的。

当然,动态配置在工程上如何更合理地实现,也是一个值得探讨的话题,这个留待以后再说。

实验分组信息记录

做 ABTest 的目的,主要是为了收集用户对实验的反馈,方法则是查看不同分组下用户数据指标和用户行为的对比和变化。这些数据的记录,往往是靠各种系统日志。

一种比较原始的办法,是用系统日志跟用户中实验分组的时间信息去 join,来区分 AB 分组。比如用户 A 在 1 号到 10 号分到了 A 分组,那么可以认为该用户 1 号到 10 号的日志都属于 A 分组,用来统计 A 分组指标。

这种方式存在着明显的缺点:一是实验的开启和关闭点不会是整点,而进行精确到秒的日志时间 join 成本太高,很多时候不得不抛弃首尾两天不足整天的用户日志;二是用户中实验分组在产品上的生效往往不是实时的。尤其是在客户端上,用户正在用一个功能的时候,很难瞬间将该功能切换成另一种样式,往往是在用户在下一次重入初始化的时候才开启实验样式,否则很容易引起崩溃。

对统计更友好的分组信息记录,是在每条关键的系统/客户端日志中都添加分组信息字段。这样虽然增加了一些冗余信息,但会使所有的关键数据记录都有实验分组这一列,用它做筛选即可进行指标的对比和用户行为序列的分析。

小细节:分组编码

上面谈到每条日志中都要记录用户所属的实验分组信息,这种冗余程度要求分组信息有着非常集约的编码设计,才能尽量减少传输和存储的数据量。可能的选择有以下几种:

  • 奔放派 :直接使用实验名+分组名作为分组信息。比如一个用户中了两个实验的不同分组,那就是:“ExpTitle_GroupA|ExpImage_GroupC”。奔放派的好处是日志可读性很高。
  • 婉约派:将实验名和/或分组名精简为两个 ID,形如:“10010_1|10080_3”。好处是虽然可读性低了些,但毕竟直观。
  • 理工派:将实验名和分组名精简为一个ID,形如:“12345|14523”。对于理工科思维来说,只要 ID 不重复就行,为啥要分成俩字段?
  • 极客派:将 ID 用 62 进制表示,形如:“3d7|3Mf”。这个世界上,只有麻瓜才用 10 进制。
  • 抠门派:将 ID 数组用 protobuf (varint)表示,有时候需要 base64 一下,形如:“算了举例太费劲”。好处就是一般人看不懂。

当然,可能还有其它的组合,大家意会就好了。

下篇,我们聊一下灰度发布和早鸟用户

ABTest 平台设计 - 如何进行流量分桶

在 2018 年,我相信 ABTest 这个名词已经不用过多地解释了。但我发现很多公司,尤其是初创企业,虽然能理解这件事是什么,却不知道这件事该怎么做,以及该怎么做好。

这一系列文章,就是想讲清楚在设计具体的 A/B 测试平台这种基础架构时,要考虑哪些问题,以及有哪些推荐的做法。

今天先谈一谈:

如何进行用户分桶

我们都知道互联网产品的 ABTest 主要是围绕用户进行的实验,从统计意义上观察用户对不同的产品设计、交互体验、业务流程的反馈,从而指导产品的改进方向。

那么很重要的一点就是,怎么圈定哪些用户进行 A 实验,哪些用户进行 B 实验。

一种错误做法

在我工作过的一家公司,他们是这样做的:

“使用用户的 UserID 对 1000 取模分成 1000 个桶,然后选择不同的桶分配给 A 或者 B。”

我问研发人员为什么这么做?他们给的理由是:

“UserID 是自增 ID,跟用户注册顺序有关,有一定的随机性。可以保证用户随机地分到 A 或者 B 中。”

A/B 的流量圈定的一个重要原则就是无偏,不然无法进行对比评估。上面的做法看起来倒也有一定的道理。(还常见的一种做法是,用手机尾号最后一位来进行分桶,优惠多少就看你手机尾号是否运气好了 ^_^ )

单单考虑孤立实验,这样做也无可厚非。但如果考虑到长期交叉、连续的实验,这样做有很大的问题。

首先,这种设计只能进行单层实验,也就是说一份流量只能通过一个实验。

如果实验人员选择了在任意一个桶中同时进行 X, Y 两个实验的话,那两个实验的结果就会相互干涉,导致最终结果不可信。例如:在尾号为 001 的桶里进行了两个促销活动“降价10%”和“满100减10块”的实验,最终 001 桶的用户订单数比其它桶高,那到底是哪个促销更有效果呢?

其次,这种设计在长期会造成桶间用户行为有偏

也许刚开始因为其随机性,桶间用户行为差异很小。但第一个实验过后,桶间就开始有了行为差异——这也是 ABTest 的目标。N 个实验过后,桶间行为的差异可能就变得非常大了。

比如你总是在 001 桶的用户上实验幅度较大的促销活动,那么 001 桶的用户留存就会显著高于其它桶。那实验人员为了让实验效果更好看,可能会偷偷地继续选择 001 桶进行实验。

最后,这种设计的实验效率太低。因为一份流量只能通过一个实验,无法对流量进行充分的利用。

那该如何设计用户分桶,才能满足 ABTest 的需求呢?

一种正确方法

目前业界应用最多的,是可重叠分层分桶方法。

具体来说,就是将流量分成可重叠的多个层。因为很多类实验从修改的系统参数到观察的产品指标都是不相关的,完全可以将实验分成互相独立的多个层。例如 UI 层、推荐算法层、广告算法层,或者开屏、首页、购物车、结算页等。

单单分层还不够,在每个层中需要使用不同的随机分桶算法,保证流量在不同层中是正交的。也就是说,一个用户在每个层中应该分到哪个桶里,是独立不相关的。具体来说,在上一层 001 桶的所有用户,理论上应该均匀地随机分布在下一层的 1000 个桶中。

通过可重叠的分层分桶方法,一份流量通过 N 个层可以同时中 N 个实验,而且实验之间相互不干扰,能显著提升流量利用率。

从实操上来说,我们通常采取下面的方法:

首先,确定 Layer,确定请求 Tag。例如从 UserID,DeviceID、CookieID、手机号 中选一个,支持匿名流量的,一般会选用 DeviceID 或者 IMSI 等作为请求 Tag。

然后,选一个你喜欢的 Hash 函数,尽量选个使用方便、随机性更强的;

最后,通过 Hash(Layer, Tag) % 1000 确定每层分桶。如果 Hash 函数支持 seed,那么使用 Layer 作为 seed,否则作为 SALT,即将 "Layer+Tag" 作为输入参数。

在完成分桶以后,还可以进行一定的流量筛选。例如来自北京和上海的用户,可以允许分别进行不同的实验。

可重叠分层分桶方法的系统性介绍,可以参见 Google 在 KDD 2010 发表的论文 《Overlapping Experiment Infrastructure: More, Better, Faster Experimentation》,感兴趣的同学可以延伸阅读一下。

更多样的选择

但分层方法并不能让所有人满意,尤其是对并行实验非常多的大公司来说。因为同一层可能有不同份额的其它小流量实验在线上,新实验能够拿到的流量非常有限,需要等待同层其它实验结束,会非常影响迭代效率。而新实验未必一定会修改其它实验的参数,或者影响其它实验的效果。

所以一些公司也逐渐开始探索无分层,也可以称为无限分层的方法。具体的做法是,每个实验都可以看成独立的层,只保证层间流量分桶的正交性,所有的实验都可能存在重叠情况。

这样做用户就有大概率同时中同一类实验,例如两个实验人员同时在实验 UI 布局,实际上用户只能看到一种布局(因为代码分支有先后逻辑),那么实验结果就不可信了。

由于系统上无法保证实验间不冲突,那么只能从组织上来避免冲突,或者从数据上尽早观察到冲突的存在来解决冲突。这对组织的管理能力和企业的数据能力提出了更高的要求。

下篇,我们聊一下实验的开关和分组信息传递

ETH ICAP 地址协议算法实现

以太坊钱包生成收款码时,有些是直接拿裸地址例如 "0x0728F0...75445F" 生成收款二维码,例如 TrustWallet;还有些是使用 "iban:" 开头的 ICAP 串来生成收款二维码,例如 imToken 1.0。

IBAN 本身是国际上一部分银行间转账使用的账号代码格式,以太坊社区在 IBAN 地址格式上做了一些扩展,用来编码以太坊地址和校验码,用作地址交换使用,叫做 ICAP (Inter exchange Client Address Protocol)。

IBAN 编码看起来很简单,但在实现上字母到数值的转换方法挺 trick 的,需要花一些时间进行理解。为简化理解,下面我拿一个例子来说明整个编码过程。

假设我们已经有了一个以太坊地址:0x730aEA2B39AA2Cf6B24829b3D39dC9a1F9297B88,下面是生成对应 ICAP 地址的过程:

第一步:将原始 16 进制以太坊地址转换成为 36 进制地址:

16 进制 ETH 地址:0x730aEA2B39AA2Cf6B24829b3D39dC9a1F9297B88
36 进制 ETH 地址:DFRZLRUTFTFY4EVINAHYF7TQ6MACYH4

第二步:为 36 进制 ETH 地址拼接上国家码 "XE" 和空校验字符串 "00" 形成 36 进制待校验字串:

36 进制 ETH 地址: DFRZLRUTFTFY4EVINAHYF7TQ6MACYH4
36 进制待校验字串: DFRZLRUTFTFY4EVINAHYF7TQ6MACYH4XE00

第三步:将 36 进制待校验字串逐字符转成 10 进制数字字串:

36 进制待校验字串: DFRZLRUTFTFY4EVINAHYF7TQ6MACYH4XE00
10 进制待校验字串: 1315273521273029152915344143118231017341572926622101234174331400

第四步:将 10 进制大整数对 97 取模,然后用 98 - 模数:

校验码:42 = 98 - 1315273521273029152915344143118231017341572926622101234174331400 % 97

第五步:将校验码替换空校验字符串,然后重新安排 XE** 到地址前,并加上前缀:

36 进制待校验字串: DFRZLRUTFTFY4EVINAHYF7TQ6MACYH4XE00
36 进制已校验字串: DFRZLRUTFTFY4EVINAHYF7TQ6MACYH4XE42
36 进制 IBAN 号: iban:XE42DFRZLRUTFTFY4EVINAHYF7TQ6MACYH4

可以用这个 IBAN 号生成二维码,用支持 iban 地址的 app 扫描二维码验证下是否能解析到正确的 ETH 原始地址。

ICAP 地址生成和校验的实现,可以参考下面这段 Java 代码,可直接用于 Android 客户端:

package com.yangwenbo;

import java.math.BigInteger;

/**
 * Ethereum ICAP (Inter exchange Client Address Protocol) Address Converter
 * Convert Ethereum Address from/to ICAP iban address
 *
 * @ref https://github.com/ethereum/wiki/wiki/Inter-exchange-Client-Address
 * -Protocol-(ICAP)
 */
public class EthICAP {
  private static String ICAP_XE_PREFIX = "XE";
  private static String IBAN_SCHEME = "iban:";
  private static String IBAN_MOD = "97";

  /**
   * Build ICAP iban address from ethereum address.
   *
   * @param ethAddress ethereum address
   * @return ICAP iban address
   * @example input:  0x730aea2b39aa2cf6b24829b3d39dc9a1f9297b88
   * return: iban:XE42DFRZLRUTFTFY4EVINAHYF7TQ6MACYH4
   */
  public static String buildICAP(String ethAddress) {
    if (!ethAddress.startsWith("0x") || ethAddress.length() != 42) {
      throw new IllegalArgumentException("Invalid ethereum address.");
    }
    BigInteger ethInt = new BigInteger(ethAddress.substring(2), 16);
    String base36Addr = ethInt.toString(36).toUpperCase();
    String checkAddr = base36Addr + ICAP_XE_PREFIX + "00";
    String base10Str = "";
    for (Character c : checkAddr.toCharArray()) {
      base10Str += new BigInteger(c.toString(), 36);
    }
    Integer checkSum = 98
        - (new BigInteger(base10Str)).mod(new BigInteger(IBAN_MOD)).intValue();
    String icapAddress = IBAN_SCHEME + ICAP_XE_PREFIX 
        + checkSum.toString() + base36Addr;
    return icapAddress;
  }

  /**
   * Decode ethereum address from ICAP iban address
   *
   * @param icapAddress ICAP iban address
   * @return ethereum address
   * @example input:  iban:XE42DFRZLRUTFTFY4EVINAHYF7TQ6MACYH4
   * return: 0x730aea2b39aa2cf6b24829b3d39dc9a1f9297b88
   */
  public static String decodeICAP(String icapAddress) {
    if (!isValid(icapAddress)) {
      throw new IllegalArgumentException("Invalid icap address.");
    }
    BigInteger ethInt = new BigInteger(icapAddress.substring(9), 36);
    String base16Addr = ethInt.toString(16).toLowerCase();
    return "0x" + base16Addr;
  }

  /**
   * Check ICAP iban address validation
   *
   * @param icapAddress ICAP iban address
   * @return true if valid; false if invalid
   */
  public static boolean isValid(String icapAddress) {
    if (!icapAddress.startsWith("iban:XE") || icapAddress.length() != 40) {
      return false;
    }
    String base10Str = "";
    for (Character c : icapAddress.substring(9).toCharArray()) {
      base10Str += new BigInteger(c.toString(), 36);
    }
    for (Character c : icapAddress.substring(5, 9).toCharArray()) {
      base10Str += new BigInteger(c.toString(), 36);
    }
    Integer checkSum
        = (new BigInteger(base10Str)).mod(new BigInteger(IBAN_MOD)).intValue();
    return checkSum == 1;
  }
}

专门用于批量空投的 ETH 智能合约

2018 年 4 月份,美链 BeautyChain (BEC) 爆出了智能合约漏洞,导致市值 60 亿的代币价格几乎归零,这在币圈里也算是一场大事件了。

那么漏洞主要出在什么地方呢?主要是 BEC 在标准的 ERC20 接口之外,自己添加了一个 batchTransfer 接口。一般而言,添加这种接口主要是为了便于空投时批量转账,可是 BEC 在这个接口设计上犯了两个错误:

  1. 没有使用安全的数值计算,存在数值溢出漏洞。这也是在网上被广泛传播的漏洞分析。
  2. 没有限定 batchTransfer 的使用范围。如果它限制了 batchTransfer 只能合约拥有者调用,即使存在漏洞也不能被人利用了。

BEC 为了便于空投增加的这个接口可谓代价惨重,但如果他们知道还有不需要修改代币主合约就能批量转账的办法,不知道会不会吐一口老血?

批量转账,指的是在一笔 ETH 交易中转多笔代币到不同的账户,一般用于 ERC20 代币项目启动时对用户进行空投(有人叫糖果发放)。

批量空投的好处主要有两个,一是省 GAS 费,但事实上省得不多;二是省时间,这是最主要目的。以太坊是以交易为粒度打包,如果转账只能单对单,即使一次发起多笔单对单的交易,等待这些交易被打包的时间也非常漫长,而且还有笔数上限限制。将多笔转账放到同一个交易中,被打包确认的速度就会非常快。一般 ERC20 代币项目启动时都会大撒币,空投地址动辄都是几万几十万,批量空投接口对效率会有上百倍的提升。

空投合约基本原理:ERC20 可以通过 approve 和 transferFrom 两个接口授权其它地址一定的额度。那既然是这样,我们也可以授权一个合约地址来花自己的代币,如果这个合约支持批量转账,那么就可以通过这个合约来实现批量空投了。

下面是具体实操流程:

假设已经存在一个 ERC20 代币的合约,合约地址为“TOKEN_ADDR”,而你的钱包里已经有了 100 万 TOKEN,你的钱包地址是“WALLET_ADDR”。

STEP1: 用自己的钱包部署支持批量转账的空投合约,假设创建成功后地址为“AD_ADDR”。下面给出最关键的部分,完整合约参考 github 链接 。这里 transferFrom 取的是 msg.sender,理论上来讲不加 onlyOwner 限定这个合约也可以给其它人使用,但为了安全还是加上较为稳妥。

contract Airdropper is Ownable {
    function multisend(address _tokenAddr, address[] dests,
                       uint256[] values) public onlyOwner returns (uint256) {
        uint256 i = 0;
        while (i < dests.length) {
           ERC20(_tokenAddr).transferFrom(msg.sender, dests[i], values[i]);
           i += 1;
        }
        return(i);
    }
}

STEP2: 用自己的钱包授权空投合约地址 AD_ADDR 100 万 TOKEN 额度。即通过代码或者 remix 执行 approve(AD_ADDR, 1000000*精度),注意这是 ERC20 合约里的接口,需要将交易发往 TOKEN_ADDR。

STEP3: 检查 AD_ADDR 是否得到了 100 万 TOKEN 授权。通过代码或者 remix 执行 allowance(WALLET_ADDR, AD_ADDR),如果结果是 100 万 TOKEN,说明空投合约已经得到你的 100万额度授权。

STEP4: 用自己的钱包调用空投合约的 multisend 接口发起批量空投。通过代码或者 remix 执行 multisend(TOKEN_ADDR, [addr1, addr2, ...], [value1, value2, ...]),执行成功即能实现批量转账。这是空投合约里的接口,需要将交易发往 AD_ADDR。

这里最容易混淆的是几个地址,TOKEN_ADDR/WALLET_ADDR/AD_ADDR,在每一步操作中要想明白调用的是哪个合约的接口,参数应该填哪个地址。 第一步只需要操作一次,第四步可以操作很多次,第二步可根据需求随时调整授权额度。

有了这个空投合约,你就可以将自己钱包里任意类型代币都通过批量转账方式空投出去。额度可控,任意账户均可用,还避免了在代币主合约里额外增加非 ERC20 标准接口带来的风险。

为什么两笔 Token 转账消耗 GAS 不同

大家都知道,以太坊 Token 转账的过程就是智能合约执行的过程,所以每笔交易都会根据智能合约的执行情况消耗一定数量的 GAS 作为交易手续费,同时也限制了交易对资源的滥用。

大家都能接受交易有手续费的概念,因为银行转账,或者支付宝微信转账,也都有可能产生手续费。大部分情况下,手续费是按照金额一定比例收取的。但会出乎很多人意料的是,在以太坊上即使往同一个地址里转账同一种 Token,消耗的 GAS 也有可能不同。

随便在 etherscan.io 上找一个近期有多笔 Token 交易的地址(其实找了好一会儿),比如这个 : https://etherscan.io/address/0x5debb351b536eb1a61be12810abe614485167f46#tokentxns

可以看到近期正好有两笔 Rating Token 转入到这个账户,对比下 GAS 的消耗:

可以看到两笔转账消耗的 GAS 分别是 37434 和 22434,差了 15000。这就奇怪了,第一笔转账的 Token 金额比第二笔少,但是消耗的 GAS 反而更多,这完全不合理啊!

其实这种现象的解释也很简单,通过对比两笔交易的 VMTRACE 可以发现,最大的差异出现在一条指令上:“SSTORE”。第一笔交易的 SSTORE 消耗了 20000 GAS,第二笔交易只消耗了 5000 GAS。

这时候只好去查文档了,以太坊黄皮书 https://ethereum.github.io/yellowpaper/paper.pdf 附录G:FEE SCHEDULE。发现下面这段说明:

当 SSTORE 将存储的值从 0 改成非 0 时,消耗 20000 GAS,其它情况下消耗 5000 GAS。真相大白,不是以太坊乱收费,文档就是这么写的。

虽然明白了原理,但这的确颠覆了我们的认知,转账的 GAS 费用居然跟对方的账户余额有关!

不过这也解释了一件事:为什么我们在 etherscan 上经常能看到那么多 *.9999999 的转账?可能很多交易所或者黑客早就弄明白了这件事,故意留一点点金额在账户里,减少未来转账的 GAS 手续费。

JS Bridge API - 安卓和iOS统一设计探讨

1 背景

APP 开发过程中,为了追求开发效率、更新成本、性能和交互体验的平衡,经常会采取 Hybrid 的 APP 端架构。用基于 HTML5 的 WEB APP 实现易变的业务部分,用原生代码实现对效率、权限、数据交换等有要求的功能部分,然后通过 JS Bridge 打通两者,实现 JS 与 原生代码的相互调用,完成整个产品功能。

但谈到 APP 开发,大家都知道至少存在两个平台,那就是 Android 和 iOS。这两个系统采取不同的原生开发语言,也有不同的 Webview 浏览器环境。但 WEB APP 是跨平台的,所以跨浏览器的调用总归需要在一个层面上得到统一,这样才不需要专门针对两个平台开发不同的 WEB APP。

下面先对在目前的技术框架下有哪些 JS - NA 相互调用方式做一下综合介绍,然后基于上述技术提出几种跨平台 JS Bridge API 统一设计思路,最后扩展讨论下 JS Bridge 设计中的一些值得注意的点。

2 在原生代码中调用 JS 代码

2.1 Android Platform

loadUrl 方法

Android Webview 的 loadUrl 接口,可以直接在 Java 代码中执行 Javascript 脚本。在 API 23(Android 6.0)及之前,这里的 Javascript 脚本能够获取当前加载页面的变量,甚至执行当前加载页面里定义好的函数。也就是说,传入的 JS 脚本是在当前加载页面的上下文中执行的。

    // Javascript: 在网页内定义一个函数,显示一个消息,返回一个内容
    function propose(msg) {
      alert(msg);
      return "Yes!";
    }

    // Java: 执行当前加载页面中定义好的一个函数 propose()
    webView.loadUrl("javascript:propose('Will you merry me?');");

可惜的是,这种方法:

  • 只能执行 JS,无法获取返回结果,需要用其它的方式(下文介绍)获取返回结果;
  • 而且会触发一次页面的刷新,可能会导致焦点丢失,软键盘消失之类的问题;
  • 在 Android 7.0 以后,存在兼容性问题;

evaluateJavascript 方法

不过,如果 APP 适配的版本在 API 19(Android 4.4)以后,也可以使用 Webview 的 evaluateJavascript 接口 。这也是更为推荐的做法,因为避免了上面 loadUrl 的问题。

    // Java: 执行当前加载页面中定义好的一个函数 propose()
    webView.evaluateJavascript("propose('Will you merry me?')", new ValueCallback() {
      @Override
      public void onReceiveValue(String answer) {
        // 拿到 answer 是 "Yes!"
      }
    });

间接方法:Web Event 分发

这种方法很少有人提到,因为它是一种间接的调用方法。Web Event 接口提供了一种在 DOM 里进行广播的机制,那也就意味着原生代码可以不知晓 JS 的函数名,而只是广播一个事件,由页面内的 JS 决定是否处理这个 Event。这能够避免 JS 代码执行的异常,更常用于原生代码主动通知页面某些信息更新的场景。

    // Javascript: 在网页内定义一个函数,显示一个消息,返回一个内容
    function propose(e) {
      alert(e.msg);
      return "Yes!";
    }
    // 注册 WebDomEvent handler
    window.addEventListener("propose_event", propose);
    // Java: 
    webView.evaluateJavascript("var e=new Event('propose_event'); e.msg='Will you merry me?'; window.dispatchEvent(e);", new ValueCallback() {
      @Override
      public void onReceiveValue(String answer) {
        // nothing
      }
    });

这种方法也存在无法获取返回结果的问题,也需要用其它的方式(下文介绍)获取返回结果。不过在使用到 Event 通知的场景下,我们一般也不需要返回。

2.2 iOS Platform

讲到 iOS,必须提到两个不同的 WebView,一个是过时但广泛使用的 UIWebView,另一个是建议且逐渐流行的 WKWebView。

UIWebView: stringByEvaluatingJavaScriptFromString 方法

UIWebView 提供了 stringByEvaluatingJavaScriptFromString 接口 ,并且能够获得返回结果。

    // OC: 执行当前加载页面中定义好的一个函数 propose()
    [_webView stringByEvaluatingJavaScriptFromString:@"propose('Will you merry me?')"];

这个方法的主要问题在于,它是一个同步的方法。它可能会阻塞 UI 线程,不太适合执行复杂的调用。

UIWebView: JavaScriptCore

在 iOS 7 之后,苹果提供了一个获取 UIWebView 中 JSContext 的方法,直接将 JS 执行环境暴露给原生代码。这样就可以在原生代码中任意执行 JS 代码了。同时,这个接口也可以用于 JS 调用原生代码的能力,下文中会介绍。

    // OC: 获取 JSContext 
    JSContext *context = [_webView valueForKeyPath:@"documentView.webView.mainFrame.javaScriptContext"]
    [context evaluateScript:@"propose('Will you merry me?')"];

WKWebView: evaluateJavaScript 方法

可以看到,JavaScriptCore 使用起来极其方便,但在 WKWebView 中我们享受不到这种方便了。因为 WKWebView 的页面渲染是在独立的进程中,在当前进程无法直接拿到 JSContext。

不过 WKWebView 提供了一个更好的 evaluateJavaScript 接口 ,可以传入一个回调函数,实现了 JS 的异步调用。

    // OC: 执行当前加载页面中定义好的一个函数 propose() 
    [_webView evaluateJavaScript:@"propose('Will you merry me?')" completionHandler:^(id _Nullable result, NSError * _Nullable error) {
      // 拿到 result 是 "Yes!", error 是 nil
    }];

可以看到,evaluateJavaScript 接口与上文 Android evaluateJavaScript 接口极为类似。

间接方法:Web Event 分发

当然,由于 Event 接口是 WEB 标准,iOS 上也可以同样进行 Event 分发。场景和作用请看上文,不再赘述,简单代码如下:

    // Javascript: 在网页内定义一个函数,显示一个消息,返回一个内容
    function propose(e) {
      alert(e.msg);
      return "Yes!";
    }
    // 注册 WebDomEvent handler
    window.addEventListener("propose_event", propose);
    // OC: 执行当前加载页面中定义好的一个函数 propose() 
    [_webView evaluateJavaScript:@"var e=new Event('propose_event'); e.msg='Will you merry me?'; window.dispatchEvent(e);" completionHandler:^(id _Nullable result, NSError * _Nullable error) {
      // nothing
    }];

3 在 JS 代码中调用原生代码

3.1 Android Platform

addJavascriptInterface 方法

Android 从 API 1 就开始提供了 addJavascriptInterface 接口,用这个接口可以很方便地把原生的方法注入到 JS 上下文中,可以说比 iOS 做得好很多。

    // Java: 定义一个类,提供一个接口,返回一个内容
    class NativeApis {
      @JavascriptInterface
      public String propose(String msg) {
        return "Yes!";
      }
    }
    webView.addJavascriptInterface(new NativeApis(), "Bridge");
    // Javascript: 执行一个 native 的方法
    alert(window.Bridge.propose("Will you merry me?"));

但问题在于在 API 17 (Android 4.2) 之前这个方法存在安全漏洞,攻击者可以执行任意代码。在 API 17 及以后,通过显式地给出 @JavascriptInterface 限定暴露的接口,避免了安全漏洞。但在 API 17 以前,不建议使用此方法,可以考虑下面的 work around。

URL 拦截:shouldOverrideUrlLoading

这是一种曲线救国的方式,那就是通过加载非标准 Scheme( 非 http/s, 非 ftp)的 URL,用一个非法(或者叫自定义)的 URL 传递参数。当页面中的 Javascript 动态插入一个 iframe 元素时,iframe 的 url 会被 WebView 通过 shouldOverrideUrlLoading 方法传给 WebViewClient 判断是否需要加载该 URL。在这里可以拦截自定义的 URL Scheme,通过 encode 到 URL 中的信息传递参数。

    // Java: 解析要加载的 URL,如果是自定义 scheme,调用相关的函数
    class MyWebViewClient extends WebViewClient {
      @Override
      public boolean shouldOverrideUrlLoading(WebView view, String url) {
        if (url.startsWith("bridge://")) {
          // 解析 // 后面的 action 和参数,调用相关的函数
        }
        return true;
      }
    }
    webView.setWebViewClient(new MyWebViewClient());
    // Javascript: 用不可见 iframe 打开一个自定义 URL,参数需要 urlencode
    bridgeFrame = document.createElement('iframe');
    bridgeFrame.style.display = 'none';
    bridgeFrame.src = 'bridge://propose?msg=Will%20you%20merry%20me%3F';
    document.documentElement.appendChild(bridgeFrame);

URL 拦截的问题也是无法直接拿到原生代码的返回结果,需要用 URL 字符串参数传入一个回调函数,然后用上文讲到的原生代码调用 JS 的方式回调传回结果。

弹出框拦截

Android Webview 可以定义一些接口,重载 onJsAlert()、onJsConfirm()、onJsPrompt() 这些回调方法。当 JS 控制弹出框时,这些回调会被调用,进而可以通过约定的特殊内容格式判断是真正的弹出框,还是 JS 到 NA 的调用。由于 onJsPrompt 可以返回结果,所以更合适一些。

    // Java: 重载 onJsPrompt 方法,提取 prompt 内容判断是否需要拦截
    class MyWebViewClient extends WebChromeClient {
      @Override
      public boolean onJsPrompt(WebView view, String url, String message, String defaultValue, JsPromptResult result) {
        if (message.startsWith("bridge://")) {
          // 解析 // 后面的 action 和参数,调用相关的函数
          result.confirm("Yes!");
        }
        return true;
      }
    }
    webView.setWebViewClient(new MyWebViewClient());
    // Javascript: 调用 prompt 弹框,使用特定内容格式以利于拦截
    alert(window.prompt('bridge://propose?msg=Will%20you%20merry%20me%3F'));

Local Server

APP 可以在手机的本地地址 127.0.0.1 上启动一个 HTTP/WebSocket 服务,浏览器内的 JS 可以通过本地回环网络连接到这个服务,把 APP 视为一个服务端,进行正常的 B/S 通信,也可以实现在 JS 中调用原生代码。

使用这种方式时,额外注意一点是要进行有效地鉴权。因为除了 APP 内的 WebView,手机内其它的 APP 也可以访问这个服务,很可能会造成一些安全问题。所以可能需要 NA 在加载 Webview 的时候,通过 Cookie/URL参数/JS 上下文环境传入合法的 Token,才能保证其安全性。

还有一点是,如果不幸出现了端口冲突,需要有办法去解决。

3.2 iOS Platform

URL 拦截:shouldStartLoadWithRequest

UIWebView 原生并没有提供任何可以在 JS 代码中调用 NA 方法的 API,但 UIWebView 也可以通过与 Android 相同的方式进行 URL 拦截,进而间接实现 JS 到 NA 的传参。

    // UIWebView
    - (BOOL)webView:(UIWebView *)webView 
    shouldStartLoadWithRequest:(NSURLRequest *)request 
     navigationType:(UIWebViewNavigationType)navigationType;

这个方式在 WKWebView 上,依然有效,只是叫做 decidePolicyForNavigationAction

    - (void)webView:(WKWebView *)webView 
    decidePolicyForNavigationAction:(WKNavigationAction *)navigationAction 
    decisionHandler:(void (^)(WKNavigationActionPolicy))decisionHandler;

UIWebview: JavaScriptCore

大概苹果官方也觉得这种方式太 ugly,所以后来在 iOS 7 以后,提供了一个好一些的接口,就是 JavaScriptCore。在页面加载完后,可以获取当前加载页面的 JavaScript 上下文执行环境 JSContext。然后可以把一些原生方法注入到 JSConext 中,这样页面内的 JS 就可以直接调用到这些注入的方法了。

    // OC: 获取 JSContext,将原生方法注入进去
    JSContext *context = [webView valueForKeyPath:@"documentView.webView.mainFrame.javaScriptContext"];
    context[@"propose"] = ^(msg) {
      return @"Yes!";
    };
    // Javascript: 调用 prompt 弹框,使用特定内容格式以利于拦截
    alert(window.propose('Will you merry me?'));

WKWebView: WKScriptMessageHandler 方法

然后到了 WKWebView,JSContext 不好使了。不过 WKWebView 提供了另外一个方法,那就是 WKScriptMessageHandler。在创建一个 WKWebView 的时候,可以通过配置将一个 WKScriptMessageHandler 对象指针和 NAME 传进去。这样在加载页面中,通过 window.webkit.messageHandlers.NAME.postMessage 就可以将消息传给原生的 WKScriptMessageHandler 对象。

    // OC: 编写 Message 回调,并注册 Message Handler
    @interface Brige : NSObject 
    - (void)userContentController:(WKUserContentController *)userContentController
          didReceiveScriptMessage:(WKScriptMessage *)message {
      if ([message.name isEqualToString:@"Bridge"]) {
        // 处理 message
      }
    }
    ...
    _bridge = [[Brige alloc] init];
    [[_webView configuration].userContentController addScriptMessageHandler:_bridge name:@"Bridge"];
    // Javascript: 发消息给注入的 Message Handler
    window.webkit.messageHandlers.Bridge.postMessage("Will you merry me?");

WKScriptMessageHandler 同样也是无法直接返回结果。

WKWebView: 弹出框拦截

与 Android 类似,WKWebView 也提供了弹出框的回调函数,可以通过此类函数实现参数的传递。

    // WKUIDelegate
    - (void)webView:(WKWebView *)webView 
    runJavaScriptAlertPanelWithMessage:(NSString *)message 
    initiatedByFrame:(WKFrameInfo *)frame 
    completionHandler:(void (^)(void))completionHandler;
    
    - (void)webView:(WKWebView *)webView 
    runJavaScriptConfirmPanelWithMessage:(NSString *)message 
    initiatedByFrame:(WKFrameInfo *)frame 
    completionHandler:(void (^)(BOOL result))completionHandler;
    
    - (void)webView:(WKWebView *)webView 
    runJavaScriptTextInputPanelWithPrompt:(NSString *)prompt 
        defaultText:(NSString *)defaultText 
    initiatedByFrame:(WKFrameInfo *)frame 
    completionHandler:(void (^)(NSString *result))completionHandler;

Local Server

见上文中对 Android Local Server 调用方式的讨论。

4 notify-fetch-run 间接机制

上文中讲到的很多还是较为直接的 JS-NA 相互调用方法,其实还有一些更开脑洞的方法。比如 notify-fetch-run 机制,不需要直接传递参数或者代码,只需要传递一个信号,然后通过可以共同访问的第三方传递真正的参数,进行执行。

4.1 notify

如果仅仅把相互调用简化成一个 0/1 信号,那除了上面讲到的内容,还有太多东西可以做为信号。比如 event,比如通过远程服务器通知之类,下面讲一个比较奇葩的方法:

notify 中的奇葩:online/offline event

HTML5 中有一对标准的 event,叫做 online/offline,可以反应当前浏览器的联网状况。而 WebView 呢,可以通过 webView.setNetworkAvailable() 来控制联网状态。那也就意味着,原生代码只要控制 webView 的联网状态变化,就可以发送 0/1 信号给 JS。JS 收到 0/1 信号后,可以通过下文 JS 调用原生的方式获取原生代码要传入的内容,然后执行这些内容。

这种方式最大的问题在于,需要非常精巧地设计整个状态流转。因为传入的信号信息量非常少,而且正常情况下网络状况的变化也会触发这两个 event。

4.2 fetch

fetch 也可以有很多种,只要是 JS 和 NA 都能访问到的目标,都可以做第三方信息交换。比如本地 socket,远端网站,或者本地文件 file://,或者 cookie,localstorage。

5 安卓 & iOS 统一 API

我们讨论 Android & iOS API 的统一,主要是在 JS 里的统一,因为只有 JS 是跨平台的。统一 API 有两种实现方法:

  • 一种是通过封装的统一,就是说 JS 与原生代码的底层通信方式是不同的,但通过一个嵌入 WebView 的 JS 库实现 API 的统一。
  • 另一种是无需封装的统一,也就是在底层通信的接口就保持了统一,在两端的 JS 代码上是完全一致的。

5.1 JS 调用原生代码

URL 拦截(Android & iOS)

从上文介绍的方法就可以直接看出,通过 URL 拦截实现 JS 调用原生代码是统一适用于所有平台的方法,而且没有版本限制。所以很多 JSBridge 都使用了这种方法以做到最大的兼容性。

    // Android Java: 解析要加载的 URL,如果是自定义 scheme,调用相关的函数
    class MyWebViewClient extends WebViewClient {
      @Override
      public boolean shouldOverrideUrlLoading(WebView view, String url) {
        Uri uri = Uri.parse(url);
        // FIXME 异常处理
        if (uri.getScheme().contentEquals("bridge")) {
          if (uri.getAuthority().contentEquals("propose")) {
            view.evaluateJavascript(uri.getQueryParameter("callback") + "('Yes!')", null);
          }
        } else {
          view.loadUrl(url);
        }
        return true;
      }
    }
    webView.setWebViewClient(new MyWebViewClient());
    // iOS OC: 解析要加载的 URL,如果是自定义 scheme,调用相关的函数
    - (BOOL)webView:(UIWebView*)webView shouldStartLoadWithRequest:(NSURLRequest*)request navigationType:(UIWebViewNavigationType)navigationType {
      NSURL * url = [request URL];
      if ([[url scheme] isEqualToString:@"bridge"] && [[url host] isEqualToString:@"propose"]) {
        NSArray *params =[url.query componentsSeparatedByString:@"&"];
        for (NSString *paramStr in params) {
          if ([paramStr hasPrefix:@"callback"]) {
            NSArray *kv = [paramStr componentsSeparatedByString:@"="];
            [webView stringByEvaluatingJavaScriptFromString:[kv[1] stringByAppendingString: @"('Yes!')"]];
          }
        }
        return NO;
      }
      return YES;
    }
    // 统一的 Javascript: 用不可见 iframe 打开一个自定义 URL,参数需要 urlencode
    bridgeFrame = document.createElement('iframe');
    bridgeFrame.style.display = 'none';
    bridgeFrame.src = 'bridge://propose?msg=Will%20you%20merry%20me%3F&callback=showResult';
    document.documentElement.appendChild(bridgeFrame);

这种方法的问题:

  1. 对 URL 格式有 UrlEncode 的要求,对于要传递复杂参数的情况不友好。比如我们需要在参数中传递一个正常的 URL,就需要对这个参数进行两次 UrlEncode,才能保证解码不出问题。
  2. 通过 iframe 打开 URL 的方式不太直观,也缺少调用成功的返回确认,需要在 JS 端再封装一下。

对象植入(Android & iOS UIWebView)

放宽兼容性限制,Android 不再兼容 4.1 及以前版本,iOS 不再兼容 iOS 6 及以前版本。那就可以直接通过 Android 的 addJavascriptInterface 和 iOS 的 JSContext 实现将要调用的方法以对象的方式注入到 JS 上下文中,同时也可以直接获得返回结果。

    // Android Java: 定义一个类,提供一个接口,返回一个内容
    class NativeApis {
      @JavascriptInterface
      public String propose(String msg) {
        return "Yes!";
      }
    };
    webView.addJavascriptInterface(new NativeApis(), "Bridge");
    // iOS OC: 定义一个类,提供一个接口,返回一个内容
    // *.h 
    #import 
    @protocol BrigeProtocol 
    - (NSString *)propose:(NSString *)msg;
    @end
    
    @interface Bridge : NSObject
    @end
    // *.m
    // 永远返回 Yes
    @implementation Bridge
    - (NSString *)propose:(NSString *)msg {
      return @"Yes!";
    }
    @end
    ...
      // 注意生命周期
      bridge = [[Bridge alloc] init];
    ...
      JSContext *context = [webView valueForKeyPath:@"documentView.webView.mainFrame.javaScriptContext"];
      context[@"Bridge"] = bridge;
    // 统一的Javascript: 执行一个 native 的方法
    showResult(window.Bridge.propose("Will you merry me?"));

对象植入(Android & iOS WKWebView)

如果使用 WKWebView,那就意味着进一步放宽了兼容性限制,因为 WKWebView 不支持 iOS 7 及以前版本。上文说到,WKWebView 不支持 JavaScriptCore,但提供了一个 WKScriptMessageHandler 方法。这也意味着我们只能将调用方式尽量往 WKWebView 的方式上统一。

WKWebView 注入的对象,只能使用 postMessage 接口,而且是注入到 window.webkit.messageHandlers 。虽然 Android 的 addJavascriptInterface 不能实现属性的注入,也就是说我们无法在原生代码中在 JS 上下文中添加一个 window.webkit.messageHandlers.NAME 这样一个对象。但我们可以在 WKWebView 中通过 addUserScript 注册一个加载页面时就执行的脚本,将 window.webkit.messageHandlers.NAME 赋给 window.NAME,就实现在注入对象层面的统一。即 Android 和 iOS 里的 Brige 对象都注入到了 window 下。

然后 Android addJavascriptInterface 注入的对象也实现一个与 WKWebView 类似的 postMessage 接口,那么两端就实现了底层接口上的统一。

    // Android Java: 定义一个类似于 WKScriptMessageHandler 的类
    class NativeApis {
      private WebView mWebView;
      public NativeApis(WebView webview) {
        mWebView = webview;
      }
      @JavascriptInterface
      public void postMessage(String msg) {
        try {
          JSONObject json_obj = new JSONObject(msg);
          final String callback = json_obj.getString("callback");
          // JS 是异步线程,转到 UI 线程执行 JS
          mWebView.post(new Runnable() {
            @Override
            public void run() {
              mWebView.evaluateJavascript( callback + "('Yes!')", null);
            }
          });
        } catch (JSONException e) {
          Log.i("Bridge", "postMessage: " + e.getMessage());
        }
      }
    };
    // 初始化 NativeApis 时多一个 webView 句柄
    webView.addJavascriptInterface(new NativeApis(webView), "Bridge");
    // iOS OC: 定义 WKScriptMessageHandler 处理接口
    - (void)userContentController:(WKUserContentController *)userContentController didReceiveScriptMessage:(WKScriptMessage *)message {
      // 解析 JSON,调用 callback 返回数据
      NSData *jsonData = [message.body dataUsingEncoding:NSUTF8StringEncoding];
      NSDictionary * msgBody = [NSJSONSerialization JSONObjectWithData:jsonData options:kNilOptions error:nil];
      NSString *callback = [msgBody objectForKey:@"callback"];
      [message.webView evaluateJavaScript: [NSString stringWithFormat:@"%@('Yes!')",
                                              callback] completionHandler:^(id _Nullable result, NSError * _Nullable error) {
        // FIXME 出错处理
      }];
    }
    ...
    [[_webView configuration].userContentController addScriptMessageHandler:self name:@"Bridge"];
    // 将 window.webkit.messageHandlers.Bridge 改名成 window.Bridge 与 Android 统一
    WKUserScript* userScript = [[WKUserScript alloc]initWithSource:@"if (typeof window.webkit != 'undefined' && typeof window.webkit.messageHandlers.Bridge != 'undefined') { window.Bridge = window.webkit.messageHandlers.Bridge;}" injectionTime:WKUserScriptInjectionTimeAtDocumentStart forMainFrameOnly:YES];
    
    [[_webView configuration].userContentController addUserScript:userScript];
    // 统一的Javascript: 给 Native 发送一个消息,通过回调返回结果
    message = {
      func: "propose",
      options : {
        msg: "Will you merry me?"
      },
      callback: showResult.name
    };
    window.Bridge.postMessage(JSON.stringify(message));

5.2 原生代码调用 JS

JS 调用原生代码,主要目的是为了增强 JS 的能力。而原生代码调用 JS 大部分情况下主要是为了便捷 JS 的调用,这可以分为三种情况:

  1. 主动设置上下文。每次加载页面必须执行一些 setup,将一些 JS 环境设置好,不需要每次都从服务器端获取。比如上文中讲到的 addUserScript 添加一个加载页面时的上下文环境。
  2. 主动发起与 JS 交互。在某些比较少见的场合下,原生代码可能想要主动将一些信息通知给 JS,尤其是一些不在官方 HTML5 支持能力的事件,比如语音的输入、扫码的结果、调用失败等等。
  3. 最常见的,是被动的回调 JS。也就是 JS 发起了一个调用,由于调用方式的限制无法返回,或者需要较长时间才能拿到结果,这就需要原生代码在执行完调用后通过回调回传给 JS。

主动设置上下文不需要 API 的统一。

主动发起与 JS 的交互场景比较少,可以有两种方法实现:一种是页面加载过程中将回调注册给 NA;另一种是通过 Web Event 的方式由 NA 广播给 JS 上下文。我们更建议通过 Web Event 的方式广播,这样不受页面加载状态之类的限制,交互上更简单。当然,也可以两种方法结合,增加一个 Event 到 NA 的注册,保证有效广播。

被动的回调 JS,实现上比较直观,只要在 JS 调用 NA 的接口中增加一个 callback 参数,NA 在完成之后回调记录下来的接口即可。

6 JS Bridge 设计上的更多考虑

6.1 是否使用第三方 JS Bridge 库

使用第三方 JS Bridge 库,理论上能避免很多烦恼,按照它的 step by step 指引,很容易就能配出来一个可以工作的 JS Bridge 环境。

但第三方库也有一些缺点。前面讲到,第三方库为了易用,往往在 NA 层和 JS 层都会做一套新的 Adapter API 封装,但不好意思的是,它提供的仍然是一套通用 API 封装,往往应用方还得在上面再封装一层业务 API。这也就意味着,每次 JS-NA 的调用,需要走下面的一套流程:

Created with Raphaël 2.2.0业务函数: proposeJS Adapter: 比如序列化参数等前面讲到的某种 Bridge 方法:比如 URL 拦截NA Adapter:比如某种路由机制原生函数

中间的三层是由第三方库实现的。如果不熟悉第三方库的代码,或者说第三方库在这三层做了过重的封装,那调试问题就会非常困难。

我上文讲到无需二次封装的统一 API,就是希望通过选取合适的 Bridge 方法,把 JS Adapter 这一层去掉或者让它尽量地薄。这样整个调用过程能得到充分地简化,更便于问题的追查和整体的设计。

第三方库还有一个问题就是,它往往追求大而全。比如有些第三方库就是想非常完整地支持 Hybrid App 的设计,但很多时候我们往往仅需要有限个接口调用而已。为了实现有限地一些功能,还得去了解第三方库的整体设计,有时候代价也高了些。

6.2 参数约束

由于 Javascript 是弱类型的语言,而 Java 和 OC 都是强类型的,在参数的互相传递时,需要进行严格的检查。虽说 addJavascriptInterface 等方法可以动态地注入无数个对象或者方法,但仍然不建议这样做,因为维护成本太高。就像 URL 拦截一样,搭桥的路有一条就足够了。

JS Bridge 的接口,就像是一个 RPC 协议。这个 RPC 协议需要有一个版本,这样我们知道哪些版本有哪些 API,更利于有效地调用。这个 RPC 协议需要约定哪些固定的字段,这样我们可以用在入口统一校验字段是否完整,字段类型是否可用。

6.3 出错信息

跨平台的接口,很多时候 DEBUG 比较困难,尤其是上文讲到一些方式无法直接返回结果,自然也无法直接返回错误。所以在接口上,要尽量考虑出错时错误信息的回传通道,例如接口需要提供出错的 callback。

那么问题来了,如果 callback 参数也写错了怎么办?总不能让 FE 看 APP 的 log 吧?

所以建议在接口设计上,增加一个全局错误的 Web Event,就像 Linux 系统下的 errno。任何 JS 调用 NA 失败或者回调失败,都通过这个 Event 分发出去,这样前端就很容易知道错在哪里了。

6.4 API 安全性

虽然网页是在 APP 自己的 WebView 中打开的,但因为网页天然具有的超链接性质,也很难保证所有可以点开的页面都是可信的,比如有些时候活动的落地页可能会到第三方页面等。所以对一些影响 APP 运行逻辑的关键 API 接口,需要做站点的白名单控制,避免第三方站点调用此类 API。

7 总结

这篇文章列举了可用于 JS Bridge 的各平台技术实现,建议了几种无需二次封装的 Android & iOS 平台 JS Bridge 统一 API 的可选方案,讨论了设计一个简洁、规范、安全的 JS Bridge API 需要考虑的问题和解决思路。希望对读者能有所助益。