PartialFC

Expanded Manuscript · Positive-Preserving Candidate Softmax扩展稿 · 保留正类的候选 Softmax

PartialFC: Training 10 Million Identities on a Single Machine

PartialFC:在单机上训练 1000 万身份

Expanded manuscriptLarge-scale classification
Xiang An, Xuhan Zhu, Yang Xiao, Lan Wu, Ming Zhang, Yuan Gao, Bin Qin, Debing Zhang, Ying Fu, Jiankang Deng

A visual walkthrough of Partial FC as positive-preserving candidate softmax: retain every target proxy, sample the negative denominator, and distribute persistent class ownership independently from per-step activation.

这个页面将 Partial FC 解释为“保留正类的候选 Softmax”:每次都保留目标代理,只采样负类分母,并把类别的长期归属与每步激活解耦。

On this page页面导航
01 · Abstract

The classifier is the bottleneck分类层才是真正的瓶颈

Large-scale face recognition may contain millions of identities. A conventional softmax classifier compares every embedding with every class center, so its weight matrix and logits grow linearly with the number of identities. Partial FC distributes the class centers across GPUs, retains every positive class, and samples only a fraction of the negatives.

大规模人脸识别可能包含数百万个身份。传统 Softmax 让每个特征与所有类别中心比较,因此分类权重和 logits 都随身份数线性增长。Partial FC 将类别中心分布到多张 GPU,始终保留全部正类,并只采样一部分负类。

On rank ii, the active set is Si=PiNiS_i=P_i\cup N_i, with a nominal budget qi=rCiq_i=\lfloor rC_i\rfloor. If the number of positive classes exceeds that budget, all positives are still retained.

在 rank ii 上,激活集合为 Si=PiNiS_i=P_i\cup N_i,名义预算为 qi=rCiq_i=\lfloor rC_i\rfloor。即使正类数量超过预算,也会保留所有正类。

02 · Introduction

Why full softmax stops scaling为什么全量 Softmax 无法继续扩展

Modern face recognition combines million-scale web data with margin-based softmax losses such as ArcFace. The embedding backbone is not the only scaling problem: the final classifier stores one center for every identity and compares every sample with every center. Its memory and compute therefore grow linearly with the class count.

现代人脸识别依赖百万级网络数据和 ArcFace 等带间隔 Softmax 损失。扩展瓶颈并不只在骨干网络:最终分类器需要为每个身份保存一个中心,并让每个样本与所有中心比较,因此显存和计算量都随类别数线性增长。

The paper argues that full inter-class interaction is not merely expensive; on noisy, long-tailed web data it can be harmful. PartialFC therefore asks a sharper question: which class centers must participate in every iteration, and which interactions can be sampled?

论文进一步指出,全量类间交互不仅昂贵,在带噪声、长尾的网络数据上还可能有害。因此 PartialFC 重新提出问题:哪些类别中心必须在每次迭代中参与,哪些交互可以被采样?

1. Inter-class conflict1. 类间冲突

The same person may be split across labels. Full softmax then pushes visually identical identities apart.

同一个人可能被拆成多个标签,全量 Softmax 会错误地把这些近似身份推远。

2. Passive tail updates2. 尾类被动更新

Rare classes seldom appear as positives but are repeatedly moved as negatives, so their centers can drift.

稀有类别很少作为正类出现,却不断作为负类被更新,类别中心容易漂移。

3. Logit explosion3. Logits 膨胀

Sharding W helps store centers, but the global batch still produces a massive B × C logit tensor.

即使分片存储 W,全局 batch 仍会产生巨大的 B × C logits 张量。

04 · Exact Model Parallel Softmax

First make full softmax exact across ranks先在多卡上精确复现全量 Softmax

One observation carries this whole section. To normalize a softmax row you do not need the row — you need two numbers about it: its maximum, and the sum of its exponentials. Both are reductions. A reduction can be assembled from partial results without any participant ever seeing the others' inputs. So if the classes are partitioned across ranks, each rank can compute its partial max and partial sum locally, and two collective calls turn those partials into the exact global values.

本节只依赖一个观察。要归一化 Softmax 的一行,其实并不需要这一行本身,只需要关于它的两个数:行最大值,以及指数和。这两个量都是规约(reduction),而规约可以由各方的局部结果拼出来,任何一方都不必看到别人的输入。因此只要类别被切分到不同 rank 上,每个 rank 就能在本地算出局部最大值和局部和,两次集合通信即可把它们变成精确的全局值。

That is why the arrangement below is exact and not an approximation. Rank ii owns the class shard WiW_i; every rank holds the same gathered global batch XRNk×dX\in\mathbb R^{Nk\times d} and multiplies out only its own vertical slice of logits Zi=XWiZ_i=XW_i. The full Nk×CNk\times C logit matrix is never assembled anywhere. Three reductions follow — max, sum, target probability — and each one moves a single scalar per sample, never one per class.

这正是下面这套流程「精确」而非「近似」的原因。rank ii 持有类别分片 WiW_i;所有 rank 拿到同一份聚合后的全局 batch XRNk×dX\in\mathbb R^{Nk\times d},各自只乘出属于自己的那一竖条 logits Zi=XWiZ_i=XW_i。完整的 Nk×CNk\times C logits 在任何地方都不会被拼出来。随后是三次规约——最大值、求和、目标概率——每次每个样本只搬一个标量,而不是每个类别一个。

01
mn=AllReducemax(maxcCiZi,nc)m_n = \operatorname{AllReduce}_{\max}\bigl(\max_{c\in\mathcal C_i} Z_{i,nc}\bigr)
Why: eZe^{Z} overflows fp16 around Z11Z\approx11. Single-device softmax fixes this by subtracting the row max; here the row is split, so each rank reduces its own local max and all ranks come away with the same mnm_n. Cost on the wire: one scalar per sample.为什么:fp16 下 Z11Z\approx11eZe^{Z} 就会溢出。单卡 Softmax 靠减去行最大值解决,而这里一行被切开了,于是每个 rank 先取本地最大值再规约,所有 rank 拿到同一个 mnm_n。通信开销:每个样本一个标量。
02
sn=AllReducesum(cCieZi,ncmn)s_n = \operatorname{AllReduce}_{\mathrm{sum}}\biggl(\sum_{c\in\mathcal C_i} e^{\,Z_{i,nc}-m_n}\biggr)
Why: the denominator is a sum over all classes, and the classes form a disjoint partition. A sum over a partition is the sum of the partial sums — so the exact global denominator falls out of local work plus one reduction. No rank ever needs another rank's logits.为什么:分母是对全部类别求和,而类别恰好构成一个不相交的划分。对划分求和 = 各部分和相加,所以精确的全局分母只需本地计算加一次规约即可得到。任何 rank 都不需要看别的 rank 的 logits。
03
Pi,nc=eZi,ncmnsn\displaystyle P_{i,nc} = \frac{e^{\,Z_{i,nc}-m_n}}{s_n}
Both global scalars are now known, so each rank finishes its own block alone. Pi is the corresponding columns of the true full softmax — not an estimate of them. The approximation Partial FC is named after arrives later, and it concerns which columns exist at all, not how they are normalized.两个全局标量都已知,每个 rank 独立算完自己那块即可。Pi 就是真实全量 Softmax 的对应列,而不是它的估计。Partial FC 之所以叫「partial」,近似发生在后面——它决定的是哪些列还存在,而不是这些列怎么归一化。
04
pn,yn=AllReducesum(pn,yn(i))p_{n,y_n} = \operatorname{AllReduce}_{\mathrm{sum}}\bigl(p^{(i)}_{n,y_n}\bigr)
Why: cross-entropy only reads one entry per row — the one at the true label. Exactly one rank owns yny_n; the rest contribute zero. So summing a mostly-zero vector of length Nk is really a gather, done with the cheaper collective.为什么:交叉熵每行只读一个元素——真实标签那一项。yny_n 只属于某一个 rank,其余 rank 贡献 0。因此对一个长度为 Nk、几乎全零的向量求和,本质上是一次 gather,只是用了更便宜的集合通信实现。
05
L=1Nknlogpn,yn\displaystyle L = -\frac{1}{Nk}\sum_{n}\log p_{n,y_n}
The result: this loss, and every gradient flowing out of it, is value-for-value what one impossibly large GPU would have computed. The whole section buys memory and compute savings for the price of three reductions of length Nk — none of which scale with C. No approximation has been made yet.结果:这个损失以及由它导出的全部梯度,与一块「无限大的 GPU」算出来的逐值相同。整节的代价只是三次长度为 Nk 的规约——都与 C 无关,换来的是显存与算力的下降。到这里为止,还没有任何近似。
(a) Full softmax single device holds all C classes 单卡保存全部 C 个类别 X Nk × d W d × C Z = XW Nk × C mn = maxc Znc sn = Σc exp(Znc − mn) Pnc = eZ−m / sn both W and Z grow with C W 与 Z 都随 C 增长 (b) Class-sharded forward exact global softmax, logits never gathered 精确的全局 Softmax,logits 从不聚合 rank 1 ··· rank k X1N × d XkN × d AllGather Nk × d values 搬运 Nk × d 个值 W1 Wk Z1 = XW1 Zk = XWk m(1) = max Z1 m(k) = max Zk AllReducemax → m · Nk scalars → m · Nk 个标量 s(1) = Σ eZ₁−m s(k) = Σ eZk−m AllReducesum → s · Nk scalars → s · Nk 个标量 Pi = exp(Zi − m) / s no collective depends on C 所有集合通信都与 C 无关 (c) Class-sharded backward local classifier update, summed feature gradients 分类器本地更新,特征梯度求和 rank 1 rank k G1 = (P1−Y1)/Nk Gk = (Pk−Yk)/Nk ∇W1 = XTG1 ∇Wk = XTGk applied in place, no communication 原地更新,无需通信 ∇X(1) = G1W1T ∇X(k) = GkWkT ReduceScattersum Nk × d gradients Nk × d 个梯度 ∇X1N × d ∇XkN × d ∇Wi never leaves its rank ∇Wi 从不离开所属 rank features and their gradients 特征与特征梯度 resident classifier shard W 常驻分类器分片 W logits and probabilities logits 与概率 collective communication 集合通信
Figure 1. Class-sharded softmax computation, redrawn from the manuscript. (a) A single device stores the full classifier W∈Rd×C and materializes all Nk×C logits, so both grow with C. (b) Each rank evaluates logits only for its resident shard Wi∈Rd×Ci; an AllReduce of local maxima followed by an AllReduce of exponentiated sums recovers the exact global normalization, so Pi equals the corresponding block of the full softmax without gathering logits. (c) Classifier gradients are computed and applied locally, while ReduceScatter sums the feature-gradient contributions and returns each rank its own N×d rows. Every collective moves at most Nk×d values and is independent of C.
图 1. 类别切分的 Softmax 计算(取自论文原图)。(a) 单卡保存完整分类器 W∈Rd×C 并实例化全部 Nk×C logits,两者都随 C 增长。(b) 每个 rank 只对本地分片 Wi∈Rd×Ci 计算 logits;先 AllReduce 局部最大值、再 AllReduce 指数和,即可精确恢复全局归一化,因此 Pi 与全量 Softmax 的对应块完全相同,且无需聚合 logits。(c) 分类器梯度在本地计算并原地更新,ReduceScatter 汇总特征梯度并把各 rank 自己的 N×d 行发回。每次集合通信最多搬运 Nk×d 个值,与 C 无关。
What is exact什么是精确的

The distributed loss and gradients are mathematically identical to single-device full softmax.

分布式损失和梯度与单卡全量 Softmax 数学等价。

What is never communicated什么从不通信

Neither the Nk×C global logits nor the d×C classifier is gathered. One step moves about 5×105 elements per rank - three orders of magnitude less.

Nk×C 全局 logits 和 d×C 分类器都不会被聚合。每步每 rank 只搬运约 5×105 个元素,比它们小三个数量级。

Why it still fails为什么仍会失败

Per-rank logit memory is 4NC bytes, independent of GPU count k. Adding ranks does not remove the global-batch logit bottleneck.

每个 rank 的 logits 显存为 4NC 字节,与 GPU 数 k 无关;加卡无法消除全局 batch 的 logits 瓶颈。

The cost reappears as storage代价从通信转移到显存

Model parallelism fixed the classifier. It did not fix the logits. Under a fixed number of classes per GPU, adding ranks pushes these two costs in opposite directions — and only one of them is the one you were trying to reduce.

模型并行解决了分类器,但没有解决 logits。在每卡类别数固定的前提下,加卡会把这两项成本推向相反的方向 —— 而其中只有一项是你原本想降下来的那个。

MW=4dC/k bytesM_W = 4dC/k \ \text{bytes}
Shrinks as 1/k. Per-GPU classifier weights — the part sharding does fix. The class dimension C is split k ways, so each rank stores a k-th of it and every GPU you add makes this term smaller.随 1/k 缩小。每卡分类器权重 —— 这是切分确实解决掉的部分。类别维度 C 被切成 k 份,每个 rank 只存其中一份,每加一张卡这一项就小一分。
Mlogit=4NC bytesM_{\mathrm{logit}} = 4NC \ \text{bytes}
Independent of k. Per-GPU logits — the part sharding does not touch. Every rank scores the whole global batch Nk against its own shard C/k, so the k cancels and adding GPUs does nothing at all.与 k 无关。每卡 logits —— 这是切分完全没碰到的部分。每个 rank 都要拿整个全局 batch Nk 去打自己那份 C/k 的分,k 被约掉,加卡毫无作用。
Σ
Mlogit/MW=Nk/dM_{\mathrm{logit}}\big/M_W = Nk\big/d
The verdict. The ratio therefore grows linearly in k. At N=64, d=512 it is already 1:1 at 8 GPUs and 10:1 at 80 — past that point you are buying classifier capacity you no longer need, in order to feed logits you cannot shrink. Figure 2 plots the totals as MFC ≈ 3MW + 2Mlogit: momentum SGD keeps three copies of W, the margin softmax two of the logits.结论。于是比值随 k 线性增长。在 N=64、d=512 下,8 卡时已是 1:1,80 卡时变成 10:1 —— 越过这一点后,你是在买已经不再需要的分类器容量,去供养一个根本压不下去的 logits。图 2 按 MFC ≈ 3MW + 2Mlogit 画总量:动量 SGD 保存三份 W,带间隔 Softmax 保存两份 logits。
(a) Communication per step classifier d × C 分类器 d × C 512M logits Nk × C logits Nk × C 512M embeddings X 特征 X 262K feature grads ∇X 特征梯度 ∇X 262K softmax stats Softmax 统计量 1K target probs 目标概率 512 103 105 107 109 elements per rank, one step 单步每 rank 搬运的元素数 communicated 已通信 never gathered 从不聚合 (b) Per-GPU FC memory 1.3 8 1M 1.8 16 2M 2.8 32 4M 4.9 64 8M 5.9 80 10M 0 2 4 6 per-GPU FC memory (GB) 每卡 FC 显存 (GB) GPUs k / classes C GPU 数 k / 类别数 C 3MW 2Mlogit (c) Logits dominate 0 20 40 60 80 100 40% at k = 8 87% at k = 80 0 25 50 75 GPUs k GPU 数 k logit share of MFC (%) logits 在 MFC 中的占比 (%) Mlogit / MW = Nk / d
Figure 2. Why exact class-sharded model parallelism still hits a memory wall (N=64, d=512, 125K classes per GPU, so C grows with k). (a) Elements moved per rank in one step, on a log axis: only embeddings, softmax statistics, target probabilities and feature gradients are communicated, while the Nk×C logits and the d×C classifier never leave their owner. (b) Per-GPU FC memory splits into a constant weight term 3MW and a logit term 2Mlogit=8NC that grows linearly with k, rising from 1.3 GB at 8 GPUs to 5.9 GB at 80 GPUs. (c) The logit share of MFC therefore climbs from 40% to 87%. Adding ranks buys classifier capacity, not logit capacity.
图 2. 为什么精确的类别切分模型并行仍会撞上显存墙(N=64、d=512、每卡 12.5 万类,因此 C 随 k 增长)。(a) 单步每 rank 搬运的元素数(对数坐标):只通信特征、Softmax 统计量、目标概率和特征梯度,Nk×C logits 与 d×C 分类器始终不离开各自的 owner。(b) 每卡 FC 显存分成恒定的权重项 3MW 和随 k 线性增长的 logits 项 2Mlogit=8NC,从 8 卡的 1.3 GB 升到 80 卡的 5.9 GB。(c) 于是 logits 在 MFC 中的占比从 40% 升至 87%:加卡只买到分类器容量,买不到 logits 容量。
05 · Partial Approximation

One training step, in five operations一个训练步,五个关键操作

The animation follows tensors rather than implementation details. Read it from features, to sharded class centers, to the sampled classifier, and finally to sparse updates.

下面的动画只追踪关键张量:从特征出发,经过类别中心分片和局部采样,最后完成稀疏更新。

Samples / embeddings样本 / 特征
Local class centers本地类别中心
Sampled negatives采样负类
Active target / update目标类别 / 更新

Phase 1: AllGather Embeddings阶段 1:AllGather 特征

Each GPU computes feature embeddings for its local batch. To compute the global loss, embeddings from all GPUs are gathered together.

每张 GPU 先计算本地 batch 的特征。为了计算全局分类损失,需要把所有 GPU 上的特征聚合到一起。

Look for: colored local batches观察:彩色本地 batch Communication: AllGather features only通信:只 AllGather 特征 Result: one global embedding matrix结果:一个全局特征矩阵

Phase 2: Class Center Partitioning阶段 2:类别中心切分

The massive classifier is partitioned by permanent ownership. In the manuscript's 8-GPU, 10M-class benchmark, each rank owns about 1.25M centers; the four-GPU animation below is schematic. Owned rows may remain in pinned host storage until activated.

巨大分类器按长期归属切分。在论文的 8 卡、1000 万类实验中,每个 rank 约拥有 125 万个中心;下方四卡动画仅为示意。所属行可以保留在锁页主存中,激活时再载入 GPU。

Look for: Wlocal shards观察:Wlocal 分片 Memory: 10M classes split across GPUs显存:1000 万类分摊到多卡 Class shards are never gathered分类器分片不会被聚合

Phase 3: PartialFC Sampling阶段 3:PartialFC 采样

Each rank keeps every locally owned positive and samples negatives without replacement to fill its nominal budget. Sorted active indices support label remapping; if positives exceed the budget, none are dropped.

每个 rank 保留其拥有的全部正类,再无放回采样负类填满名义预算。排序后的激活索引用于标签重映射;即使正类超过预算,也不会丢弃任何正类。

Look for: positive labels kept观察:正类标签始终保留 Negatives: small sampled buffer负类:小规模采样缓冲区 Skipped classes stay inactive未采样类别保持不动

Phase 4: Forward Pass阶段 4:前向计算

Normalized embeddings score only active proxies. The configured CosFace or ArcFace margin modifies target logits, then ranks reduce row maxima, denominator sums, and target probabilities for distributed sampled softmax.

归一化特征只与激活代理计算相似度。配置的 CosFace 或 ArcFace 间隔作用于目标 logit,随后各 rank 规约行最大值、分母和目标概率,完成分布式采样 Softmax。

Look for: partial logits per GPU观察:每张 GPU 的局部 logits ArcFace: target logits highlightedArcFace:目标 logit 高亮 Softmax: max, sum and target probabilitySoftmax:规约最大值、分母和目标概率

Phase 5: Backward & Update阶段 5:反向传播与更新

Sampled classifier gradients stay on the rank that owns each proxy. Feature-gradient contributions are reduce-scattered back to the original local batches; only active classifier rows receive loss gradients on this step. A sparse bank updates only those rows, while a dense optimizer may still apply momentum or weight decay to its full parameter shard.

采样分类器梯度保留在拥有对应代理的 rank;特征梯度通过 ReduceScatter 返回原始本地 batch,本步只有激活的分类器行获得损失梯度。稀疏类别库只更新这些行,而 dense optimizer 仍可能对完整参数分片施加动量或权重衰减。

Look for: prob - onehot观察:prob - onehot Synchronization: feature gradients同步:特征梯度 Loss gradients: sampled W only损失梯度:仅采样到的 W
06 · Infrastructure

Ownership, activation and storage are separate decisions类别归属、每步激活与存储位置彼此独立

Each class proxy has one permanent owner, but only the active set SiS_i enters the GPU classifier buffer. In the expanded manuscript's optional hierarchy, the complete class bank and optimizer state can live in pinned CPU memory or an NVMe-backed cache, reducing GPU-resident classifier storage from Θ(Cid)\Theta(C_i d) to Θ(rCid)\Theta(rC_i d). The released PartialFC_V2 path instead keeps each rank's complete owned shard on its GPU.

每个类别代理只有一个长期 owner,但只有激活集合 SiS_i 会进入 GPU 分类器缓冲区。在扩展稿提出的可选分级方案中,完整类别库与优化器状态可以位于锁页 CPU 内存或 NVMe 缓存中,从而把 GPU 常驻分类器存储从 Θ(Cid)\Theta(C_i d) 降至 Θ(rCid)\Theta(rC_i d)。已发布的 PartialFC_V2 路径则仍把每个 rank 拥有的完整分片放在 GPU 上。

The only subtlety is overlap between consecutive active sets: cold rows may arrive early, but repeated rows must carry the result of optimizer tt.

真正需要处理的是相邻激活集合的重叠:冷行可以提前到达,但重复行必须携带 optimizer tt 的最新结果。

Overlap-safe asynchronous hierarchical class storage pipeline A double-buffer timeline. Rows needed only by iteration t plus one are prefetched from host into slot B during iteration t. Rows shared by iterations t and t plus one are copied from slot A to slot B only after optimizer t updates them. A ready event gates Partial FC t plus one until both paths complete, while updated rows write back to host asynchronously. Build the next active set without stale rows 构建下一激活集合,同时避免旧版本行 Cold rows move immediately; overlapping rows wait for optimizer(t). 冷行立即搬运;重叠行等待 optimizer(t) 完成。 Sᵢ(t+1) = COLD ∪ OVERLAP disjoint and exhaustive 互斥且完备 PREPARE SLOT B 准备槽位 B 1 Prefetch cold rows 预取冷行 Sᵢ(t+1) ∖ Sᵢ(t) host → B · H2D 2 Patch updated overlap 补入更新后的重叠行 Sᵢ(t) ∩ Sᵢ(t+1) A → B · after optimizer(t) READY GATE 就绪门 wait for both 等待两个事件 CUDA events Slot B is readable 槽位 B 可读取 PFC(t+1) · current Sᵢ(t+1) no stale overlap rows 不存在旧版本重叠行 PIPELINED TIMELINE 流水时间线 earlier更早 later更晚 optimizer(t) done slot A槽位 A read · backward · optimizer(t) 读取 · 反向 · optimizer(t) D2H write-backD2H 写回 A remains lockedA 保持锁定 slot B槽位 B prefetch cold rows预取冷行 wait等待 patch overlap补入重叠行 B ready → PFC(t+1)
Two invariants. B becomes readable only after both the cold-row prefetch and the post-optimizer overlap patch complete. A becomes reusable only after its D2D readers and D2H write-back complete. Row versions validate host state; CUDA events enforce these dependencies.两个不变量。冷行预取与 optimizer 后的重叠行补丁都完成,B 才可读取;D2D 读取与 D2H 写回都完成,A 才可复用。行版本号用于校验主存状态,CUDA event 用于强制这些依赖。
Model parallel PFC模型并行 PFC

VMP=Θ(Nkd)+Θ(Nk)

Communication is independent of class count C.

通信量与类别数 C 无关。

Sampled data parallel采样数据并行

VDP=Θ(rCd)

Replicated sampled proxies still require gradient synchronization.

复制的采样代理仍需同步梯度。

Overlap condition重叠条件

The cold-row H2D and overlap-row handoff must both finish before the next PFC read; otherwise the ready gate stalls.

冷行 H2D 与重叠行交接都必须在下一次 PFC 读取前完成,否则 ready gate 会产生等待。

Per-rank GEMM shapes每个 rank 的 GEMM 形状

Per-rank GEMM shapes drawn to scale Three GEMM diagrams on a shared cell grid. Every block is built from discrete square cells: block width counts the classes held on the rank and block height counts the batch rows. Full model parallel is twenty cells wide, sampled data parallel sixteen cells wide but only one cell tall, and MP-PFC two cells wide by eight cells tall. Full MP Wᵀ · 512 × 125K X · 1024 × 512 Z · 1024 × 125K Full batch, C/k = 125K centers per rank. 每 rank 仍存 C/k = 125K 个中心。 The logit block dominates memory. logit 块主导显存开销。 Sampled DP Wᵀ · 512 × 100K X · 128 × 512 Z · 128 × 100K dashed = rows a gathered batch would fill 虚线 = 全局 batch 本可填满的行 Local batch only; rC = 100K replicated. 只有本地 batch;rC = 100K 逐卡复制。 Short and wide: almost no reuse per center. 又矮又宽:中心几乎无法复用。 MP-PFC Wᵀ · 512 × 12.5K Z · 1024 × 12.5K dashed = Full MP width, 10× wider 虚线 = Full MP 的宽度,宽 10 倍 X · 1024 × 512 Shard × sample: rC/k = 12.5K per rank. 分片 × 采样:每 rank 只剩 rC/k = 12.5K。 Tall and narrow: full reuse per center. 又高又窄:中心被全 batch 复用。 X features X 特征 Wᵀ classifier on this rank Wᵀ 本 rank 的分类器 Z logits Z logits capacity not used 未被利用的容量 k=8 · N=128 per rank · gathered 1024 · C=1M · r=0.1 · d=512

First-order traffic per rank每个 rank 的一阶访存量

Full MP
192.5M
Sampled DP
64.1M
MP-PFC
19.7M
050M100M150M192.5M
feature readweight readlogit write

Why identical FLOPs do not run at identical speed为什么相同的 FLOPs 跑不出相同的速度

Sampled DP and MP-PFC issue the same arithmetic, F=2NdrCF=2NdrC. Counting elements is not yet an answer either, because a kernel runs at whichever of compute or bandwidth saturates first. The quantity that decides which one is the ratio between them: how much arithmetic the GEMM performs per byte it is forced to pull from memory. For the classifier read, that ratio is set entirely by the shape.

采样数据并行与 MP-PFC 发射的算术量完全相同,都是 F=2NdrCF=2NdrC。而只数元素同样不算回答问题,因为 kernel 的速度取决于算力和带宽哪一个先饱和。决定这一点的是两者之比:GEMM 每从显存拉取一个字节,能做多少次算术。对分类器的读取而言,这个比值完全由形状决定。

Every classifier weight that a rank loads participates in exactly one multiply-accumulate per batch row it is multiplied against. So the reuse each loaded center earns is the number of rows in the GEMM. Sampled DP multiplies against the local batch and gets NN; MP-PFC multiplies against the gathered batch and gets NkNk. Same weights, same flops, k×k\times the reuse.

一个 rank 载入的每个分类器权重,与多少个 batch 行相乘,就参与多少次乘加。因此每个被载入的中心所获得的复用次数就等于 GEMM 的行数。采样数据并行只与本地 batch 相乘,得到 NN;MP-PFC 与聚合后的全局 batch 相乘,得到 NkNk。同样的权重、同样的浮点运算,复用高出 kk 倍。

Per-rank GEMM每 rank 的 GEMM Centers held持有的中心数 Rows multiplied参与相乘的行数 Reuse per center每中心复用 Classifier traffic分类器访存
Full MPC/kC/kNkNkNkNkCd/kCd/k
Sampled DPrCrCNNNNrCdrCd
MP-PFCrC/krC/kNkNkNkNkrCd/krCd/k

This is why the short, wide block is the bad one. A GEMM whose row count is small has few rows to amortize each loaded weight over, so it spends most of its time streaming the classifier past the tensor cores rather than multiplying with it — a bandwidth-bound kernel, no matter that its flop count is small. The tall, narrow block performs the same multiplications while touching a fraction of the weights, and amortizes each one over the full gathered batch, which is what moves it back into the compute-bound regime.

这正是"又矮又宽"那一块糟糕的原因。行数很小的 GEMM 没有足够的行去摊薄每个载入的权重,于是大部分时间都花在把分类器流过 tensor core,而不是拿它做乘法 —— 这是一个带宽受限的 kernel,哪怕它的浮点运算量本身很小。而"又高又窄"那一块,在只触碰一小部分权重的前提下完成了同样多的乘法,并把每个权重摊薄到整个全局 batch 上,这才把它推回算力受限的区间。

Shrinking the shard has a second, discontinuous effect. A GEMM tiles its output, and the weight tiles get revisited as the kernel walks down the rows; whether those revisits are served by cache or by HBM depends on whether the whole shard fits in L2. Sharding and sampling shrink the resident classifier by rkrk, which is what can carry it across that threshold — and crossing it changes the cost of reuse by an order of magnitude, not a few percent.

缩小分片还有第二个、且不连续的收益。GEMM 会对输出做分块,kernel 沿行方向推进时会反复访问权重分块;这些重访是由缓存还是由 HBM 承担,取决于整个分片能否放进 L2。分片与采样把常驻分类器缩小了 rkrk 倍,这个倍数足以让它跨过那道门槛 —— 而跨过与否,带来的是数量级的差异,不是百分之几。

One scoping note: all of the above concerns the forward GEMM alone. The logit block is where full model parallelism actually hurts, because it is written by the GEMM, read back for the softmax normalization and read once more in backward. Sampling shrinks that block by 1/r1/r, and it is the only term in the whole pipeline that scales with the product of batch and class count.

一点范围说明:以上讨论都只涉及前向 GEMM。全量模型并行真正吃亏的地方在 logit 矩阵 —— 它由 GEMM 写出、为 Softmax 归一化读回、反向时再读一次。采样把这一块缩小了 1/r1/r,而它是整条流水线中唯一随 batch 与类别数之积增长的项。

Re-drawn from the manuscript analysis. Sampled DP and MP-PFC have the same leading FLOPs, 2NdrC, and produce NrC logits per rank. MP-PFC wins through data movement: sharding and sampling reduce classifier traffic from Cd to rCd/k while the gathered batch reuses each active center.根据论文分析重新绘制。采样数据并行与 MP-PFC 的主导 FLOPs 都是 2NdrC,每个 rank 都产生 NrC 个 logits。MP-PFC 的优势来自数据移动:分片与采样把分类器访存从 Cd 降到 rCd/k,同时全局 batch 可以复用每个激活中心。
TDP ≈ Nd + rCd + NrC TMP ≈ Nkd + rCd/k + NrC FDP = FMP = 2NdrC
07 · Experiments and Results

Does partial sampling remain accurate?部分采样还能保持精度吗?

The paper trains on WebFace4M, WebFace12M and WebFace42M, then evaluates on LFW, CFP-FP, AgeDB, IJB-B, IJB-C and the more difficult MFR benchmark. CNNs use SGD for 20 epochs; ViTs use AdamW for 40 epochs. The central question is whether sampling improves scale and robustness without sacrificing verification accuracy.

论文使用 WebFace4M、WebFace12M 和 WebFace42M 训练,并在 LFW、CFP-FP、AgeDB、IJB-B、IJB-C 及更困难的 MFR 上评测。CNN 使用 SGD 训练 20 个 epoch,ViT 使用 AdamW 训练 40 个 epoch。核心问题是:采样能否在提升规模与鲁棒性的同时保持验证精度。

10M+ Identities1000 万+身份Scalable Training可扩展训练
10% Sampling10% 采样Preserves every positive proxy保留每个正类代理
Robust to Noise抗噪声Conflict, flips and long tails冲突、错标与长尾
r = 0.2-0.3best accuracy-efficiency region最佳精度-效率区间
19.7Mtraffic elements/rank in the paper's GEMM example论文 GEMM 示例中的每 rank 访存元素
+36.33MFR points under 40% label flips40% 错标下 MFR 提升点数
100Msynthetic classes scaled on 64 GPUs64 张 GPU 扩展的合成类别数

Datasets, protocols and training settings数据集、评测协议与训练配置

Training covers CASIA, MS1MV2, Glint360K (17M images, 360K identities), and WebFace4M/12M/42M (200K/600K/2M identities). Evaluation uses LFW, CFP-FP, AgeDB-30, CPLFW, CALFW, MegaFace, IJB-B, IJB-C and ICCV21-MFR. WebFace42M is strongly long-tailed: 44.57% of identities have fewer than ten images.

训练数据包括 CASIA、MS1MV2、Glint360K(1700 万图像、36 万身份),以及 WebFace4M/12M/42M(20 万/60 万/200 万身份)。评测覆盖 LFW、CFP-FP、AgeDB-30、CPLFW、CALFW、MegaFace、IJB-B、IJB-C 与 ICCV21-MFR。WebFace42M 长尾显著:44.57% 的身份少于 10 张图像。

Runs实验GPUBatch批量Optimizer / LR优化器 / LRSchedule训练周期
CASIA / MS1MV2 / Glint360K8× RTX2080Ti512 totalSGD / 0.132K / 180K / 600K iters
WebFace42M CNN32128 / GPUSGD / 0.420 epochs, 2 warm-up
WebFace42M ViT-384 / GPUAdamW / 10-340 epochs, 4 warm-up

Normalized margins: feature scale s=64; CosFace margin 0.4; ArcFace angular margin 0.5.

归一化间隔:特征缩放 s=64;CosFace margin 0.4;ArcFace 角度 margin 0.5。

Ablation: sampling ratio消融实验:采样率

Complete ResNet50 sampling-rate ablation. MFR columns use TAR@FAR=1e-6; Mask uses FAR=1e-4. Extremely low rates fail, while 0.1-0.3 generally recovers full FC.

完整 ResNet50 采样率消融。MFR 各列使用 TAR@FAR=1e-6,Mask 使用 FAR=1e-4。极低采样率明显退化,0.1-0.3 通常可恢复全量 FC。

Configuration配置AllAfricanCaucasianSouth AsianEast AsianMask
WF4M + FC-1.086.2583.3591.1188.1465.7972.05
WF4M + PFC-0.0474.1171.4381.7976.2552.2454.21
WF4M + PFC-0.185.7683.8291.0087.9066.0471.13
WF4M + PFC-0.386.8584.8691.5788.5767.5272.28
WF12M + FC-1.091.7090.7294.9493.4475.1080.47
WF12M + PFC-0.01387.8587.0792.3290.7068.2872.98
WF12M + PFC-0.191.2490.8094.6793.1874.9779.73
WF12M + PFC-0.291.7891.0995.0093.5375.9079.92
WF12M + PFC-0.391.8291.1495.0093.6175.5580.08
WF42M + FC-1.093.8693.3396.2095.2479.4683.90
WF42M + PFC-0.00891.2790.3495.1693.0476.9381.24
WF42M + PFC-0.193.9593.4896.3795.5180.0383.79
WF42M + PFC-0.294.0493.6796.3895.4980.0784.32
WF42M + PFC-0.394.0393.6896.3895.5279.7684.46

Source: paper Table 1. The optimum is broad rather than a single fragile setting.

来源:论文表 1。最优区间较宽,并非依赖单一脆弱超参数。

Robustness: conflict, label flips and long tails鲁棒性:类别冲突、错标与长尾

Reducing negative interaction is most valuable when labels are imperfect. PFC* additionally filters abnormal negative centers whose cosine similarity exceeds 0.4.

当标签不完美时,减少负类交互的价值最明显。PFC* 还会过滤余弦相似度高于 0.4 的异常负类中心。

Stress test压力测试FCPartialFCPFC*Best gain最佳提升
Inter-class conflict类间冲突79.9391.2091.68+11.75
40% label flips40% 标签翻转43.8778.5380.20+36.33
Long-tail identities长尾身份分布87.4491.96-+4.52

MFR-All results from Tables 3-5. Long-tail PartialFC uses r=0.2; conflict PFC/PFC* entries use r=0.1 and r=0.2 respectively.

MFR-All 结果来自表 3-5。长尾实验使用 r=0.2;冲突实验的 PFC / PFC* 分别取 r=0.1 / 0.2。

Systems scaling: class-axis sampling系统扩展:沿类别轴采样

For k=8, local batch N=128, C=1M, r=0.1 and d=512, the manuscript estimates 19.7M memory-traffic elements per rank for model-parallel Partial FC, versus 64.1M for sampled data parallelism and 192.5M for full model parallelism. Its communication scales as Θ(Nkd)+Θ(Nk), independent of C.

当 k=8、本地 batch N=128、C=100 万、r=0.1、d=512 时,论文估算模型并行 Partial FC 每个 rank 的访存量为 1970 万个元素;采样数据并行为 6410 万,全量模型并行为 1.925 亿。其通信复杂度为 Θ(Nkd)+Θ(Nk),与类别数 C 无关。

Active benchmark results当前稿件基准结果

Partial FC remains close to full FC across datasets, but does not improve every metric. MFR-All is TAR@FAR=1e-6; IJB-C is TAR@FAR=1e-4; the other columns are verification accuracy.

Partial FC 在不同数据集上总体接近全量 FC,但并非每个指标都提升。MFR-All 为 TAR@FAR=1e-6,IJB-C 为 TAR@FAR=1e-4,其余为验证准确率。

Training set训练集MFR-AllAgeDBCFP-FPLFWIJB-C
CASIA36.7999.4595.2194.9087.22
CASIA + PFC37.1199.3795.4394.6084.97
VGGFace38.5899.5597.4195.0891.22
VGGFace + PFC40.6799.6898.5395.4092.49
GlintAsian62.6699.5893.1995.4091.50
GlintAsian + PFC63.1599.6593.0395.2391.14
MS1MV277.7099.8398.0898.0896.14
MS1MV2 + PFC77.7499.7898.0798.0296.08
MegaFaceMS1M78.3799.7597.5697.4095.35
MegaFaceMS1M + PFC78.7799.8097.8797.7395.40
MS1MV382.5299.8098.5398.2796.58
MS1MV3 + PFC81.6899.8098.4498.1796.43
Glint360K86.7999.8299.1498.4597.13
Glint360K + PFC87.0899.8299.1498.4597.02
WebFace12M90.5799.8099.2098.1097.12
WebFace12M + PFC89.9599.8299.1498.1297.01

Source: active benchmark table in the expanded manuscript.

来源:扩展稿当前启用的基准表。

08 · Implementation

Minimal training loop最小训练循环

The implementation needs one global feature gather, local class sampling, distributed softmax statistics, and sparse optimizer updates.

实现只需要全局特征聚合、本地类别采样、分布式 softmax 统计,以及稀疏的优化器更新。

# Model-parallel Partial FC training step X, y = differentiable_all_gather(local_X, local_y) S_i = local_positives(y) | sample_local_negatives(r) Z_i = normalize(X) @ normalize(W_i[S_i]).T loss = distributed_sampled_margin_softmax(Z_i, remap(y, S_i)) loss.backward() # sampled classifier rows update locally; feature gradients reduce-scatter optimizer.step(sampled_rows=S_i)

Communication and memory comparison通信量与显存对比

QuantityExact model parallel精确模型并行Partial FC
Active classes / rank每 rank 激活类别C/krC/k
Classifier FLOPs分类器 FLOPs2NdC2NdrC
Communication通信复杂度Θ(Nkd)+Θ(Nk)Θ(Nkd)+Θ(Nk)
Classifier gradients分类器梯度Remain local to owner保留在所属 rankActive rows remain local激活行保留在所属 rank
Feature gradients特征梯度Reduce-scatterReduceScatterReduce-scatterReduceScatter
09 · Conclusion and Discussion

Sparse interaction is both a systems and learning choice稀疏交互既是系统选择,也是学习选择

PartialFC retains all class centers over the full training run, but selects and updates only positives plus a random subset of negatives at each step. This lowers classifier compute, reduces passive tail-class updates, and decreases exposure to conflicted negative labels. The experiments show that moderate sampling ratios preserve accuracy while improving robustness and large-class throughput.

PartialFC 在整个训练过程中保留全部类别中心,但每一步只选择并更新正类和随机负类子集。这样既降低分类器计算量,也减少尾类的被动更新和冲突负标签的影响。实验表明,中等采样率能够保持精度,同时提升鲁棒性与超大类别训练吞吐。

Approximation boundary.近似边界。 Renormalizing over the active set means the sampled gradient is generally not an unbiased full-softmax gradient. Very low sampling rates can weaken inter-class separation.在激活集合上重新归一化,意味着采样梯度通常不是全量 Softmax 梯度的无偏估计;过低采样率会削弱类间分离。

Systems boundary.系统边界。 Host or NVMe staging helps only when transfers overlap computation. Extending this implementation to dynamically encoded contrastive candidates still requires candidate communication and gradient routing through both encoders.只有传输能与计算重叠时,主存或 NVMe 分级存储才有收益。扩展到动态编码的对比候选时,仍需解决候选通信及双编码器梯度路由。

10 · References
  1. [1]
    Killing Two Birds with One Stone: Efficient and Robust Training of Face Recognition CNNs by Partial FC An, Deng, Guo, Feng, Zhu, Yang, Liu · CVPR 2022 · arXiv:2203.15565
  2. [2]
    Partial FC: Training 10 Million Identities on a Single Machine An, Zhu, Gao, Xiao, Zhao, Feng, Wu, Qin, Zhang, Zhang, Fu · ICCV Workshops 2021 · arXiv:2010.05222