alphaFM的内存优化

　　两年前把alphaFM放到了github上，之后不断有人跟我询问算法原理或反馈在使用中碰到的问题等，看起来有不少大小公司的业务有在用这个工具，我很欣慰。

　　最近浪费了几个周末，做了一次大的改动，主要是内存优化。

　　最早写这个工具没有太在意内存的使用，当把alphaFM当做LR训练工具使用时，实践发现，128G内存的机器，差不多到三四亿左右的维度时内存已经捉襟见肘。但人的欲望是无尽的，继续加特征的冲动是无法遏制的，虽然通过一些trick还可以勉强撑一阵子，终究不是正道。so，我开始重新审视代码，是时候重构了。

　　这次重构之后，不改动任何参数，在我实验中训练时内存可以降到原来的1/3左右（当然不同数据不同参数降幅可能略有不同），这样，单机支持10亿维度妥妥的~如果你是土财主，内存不是128G而是1个T，辣么，支持百亿维度也不再是梦:)

　　旧的代码发布到v1.0.0，新的代码目前在master分支。下面是具体的优化过程。

1. 锁的优化

　　占用内存主要是存储模型参数的unordered_map，key是特征名string，value是指针指向下面这个结构体。有几亿维度特征，内存里便new几亿个ftrl_model_unit：

1
2
3
4
5
6
7
8
9
10
11
12
class ftrl_model_unit
{
public:
    double wi;
    double w_ni;
    double w_zi;
    vector<double> vi;
    vector<double> v_ni;
    vector<double> v_zi;
    mutex mtx;
    ... // member functions
};

　　平常用STL用习惯了，从来没想过它们都是内存开销大户，比如sizeof(mutex) = 40, sizeof(vector<double>) = 24，第一步便拿mutex开刀。

　　每一维特征挂一个mutex其实完全没有必要，多线程训练的时候，只是那一瞬间涉及的特征需要加锁，其他特征无需保留，所以改成维护一个lock pool，用锁的时候通过特征名hash来实时查找对应的锁即可。

　　lock pool的锁数量远少于特征维度，当然会发生冲突，即多个特征对应同一个锁，冲突多了计算速度会有影响，我们可以大概估算一下维护多少个锁可以接受。

　　设有t个线程，lock pool共有m个锁，每一瞬间最多t个参数要更新，即最多申请t次锁，这t次申请至少有两次申请的锁是同一个锁的概率为：

　　P(m,t)=1-m*(m-1)*(m-2)*…*(m-t+1)/m^t

　　这里最简化了问题，认为每次申请到的锁等概率。

　　实际使用中经常设线程数为30，即t=30，那么P(100,30)=99.22%，P(1000,30)=35.55%，P(10000,30)=4.26%

　　可见lock pool维护10000个左右的锁足够，也占不了多少内存。这里我选了一个接近10000的质数10009，呃别问我为什么，我只会告诉你19，109，1009，10009都是质数，而100009，1000009，10000009都是合数…

2. 实现内存池分配

　　去掉mutex之后，class ftrl_model_unit还有碍眼的三个vector，特别是当factor_num=0时这三个vector纯属多余，白白占用3*24字节的内存。

　　首先想到的就是弃用vector，改为普通的堆上数组，ftrl_model_unit里改为三个double*指针，针对factor_num=0的特殊情况可以通过宏把这三个指针也舍弃。

　　想法是美好的，现实是打脸的，实测发现节省的内存跟预期差的远，原因在于new默认调用的是glibc的malloc，malloc再通过brk或mmap系统调用向内核申请堆内存。而glibc复杂的内存分配机制，导致当大量申请小内存时，glibc从系统实际申请的内存要比预想大得多，甚至大好多倍。

　　我又尝试了Google的tcmalloc，以及配合上用libcuckoo代替STL的unordered_map，用__gnu_cxx::__pool_alloc代替std::allocator等，有点用但都收效甚微。

　　后来想明白了，无论是malloc还是tcmalloc都是一种通用的内存管理，需要考虑分配还要考虑回收，不可能针对我这里的特殊情况做到极致优化。仔细考虑我们的需求：大量小内存分配，常驻无需中途删除，因此自己维护个内存池就可以了，每次通过malloc申请64M的大内存块，小内存就在这64M上申请，通过placement new来构建对象，64M用光了就再申请64M。维护内存池的代码非常简单，就像下面这样：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class mem_pool
{
public:
    static void* get_mem(size_t size)
    {
        if(size > blockSize) return NULL;
        if(NULL == pBegin || (char*)pBegin + size > pEnd)
        {
            pBegin = malloc(blockSize);
            pEnd = (char*)pBegin + blockSize;
        }
        void* res = pBegin;
        pBegin = (char*)pBegin + size;
        return res;
    }
private:
    mem_pool() {}
    static void* pBegin;
    static void* pEnd;
    static const size_t blockSize = 64 * 1024 * 1024;
};
void* mem_pool::pBegin = NULL;
void* mem_pool::pEnd = NULL;

　　多线程的问题，交给调用上层来解决，这里连锁都不需要。

　　注意一下：我这里实现的内存池根本没管内存对齐的问题，get_mem申请到的内存起始地址甚至可能是奇数，据说内存不对齐在某些CPU上会导致运行异常。管不了这么多，本来也不是要写全平台支持的代码，只要在Linux x86_64上能正常跑就行。如果考虑内存对齐，反而又要浪费很多内存。

3. 可变长对象

　　ftrl_model_unit中的vi，v_ni，v_zi是三个数组，长度是由参数factor_num决定的，编译期还未定，执行期才能确定数组大小。因此无法把整个数组放在ftrl_model_unit内，通常做法是ftrl_model_unit中放三个double*指针，每次先new ftrl_model_unit，然后再分别new三个数组。

　　能不能把这三个指针省掉？当然可以。每次构建ftrl_model_unit在mem_pool中get_mem时，传入的size可以不仅仅是ftrl_model_unit的大小，而是

1
size = sizeof(ftrl_model_unit) + 3 * factor_num * sizeof(double);

　　即同时给三个数组申请了内存，三个数组永远紧跟在ftrl_model_unit的屁股后面，知道了ftrl_model_unit的地址通过偏移就能访问到三个数组，无需再用三个指针专门指向它们。这样看起来就像实现了可变长的对象一样。

4. 特征名string改为char*

　　unordered_map的key是特征名字符串，在实际业务中因为经常有组合特征，导致特征名可能很长，动辄好几十个字节，吃起内存来有时比ftrl_model_unit还狠，因此也想把特征名放到mem_pool上分配。但STL的string等价于
basic_string<char, char_traits<char>, allocator<char> >，数据的内存分配由allocator<char>控制，无法指定到mem_pool上。一种方法就是实现自己的Alloc，比如：

1
using my_string = basic_string<char, char_traits<char>, my_allocator<char> >;

　　但这时的my_string和string成了两种类型，交互起来更麻烦。干脆回归原始，使用最简单的char*来保存特征名，可以很容易指定到mem_pool中分配。但当用char*作为unordered_map的key时需要自己实现hash function和key equivalence predicate，unordered_map具体如下：

1
using my_hash_map = unordered_map<const char*, ftrl_model_unit*, my_hash, my_equal>;

5. unordered_map自定义内存分配

　　unordered_map的完全体长这样：

1
2
3
4
5
6
template < class Key,                                    // unordered_map::key_type
           class T,                                      // unordered_map::mapped_type
           class Hash = hash<Key>,                       // unordered_map::hasher
           class Pred = equal_to<Key>,                   // unordered_map::key_equal
           class Alloc = allocator< pair<const Key,T> >  // unordered_map::allocator_type
           > class unordered_map;

　　乍一看以为Alloc负责pair<const Key,T>的内存分配，实际上翻看gcc的STL源码发现根本不是这么回事。

　　先回顾一下gcc的STL是如何实现unordered_map的，相关代码主要在unordered_map.h，hashtable.h，hashtable_policy.h三个头文件。

　　这里以gcc4.8.5为例，实现是标准的哈希桶的方法，class unordered_map中有一个_Hashtable类型成员变量_M_h，class _Hashtable才是具体实现。

　　_Hashtable中维护一个_Hash_node_base*指针数组_M_buckets，数组中每个指针指向一个单链表，该链表中结点包括Key哈希后落到该桶的元素，结点内存布局类似这样：

1
2
3
4
5
6
struct _Hash_node
{
    _Hash_node_base* _M_nxt;
    pair<const Key,T> _M_v;
    std::size_t _M_hash_code;
};

　　_M_hash_code缓存了Key的hash值，猜测应该是为了在rehash的时候省却重新计算Key的hash值，起到加速作用，弊端就是浪费内存。这一项在结构体中可能有可能没有，存在与否取决于Key的类型以及相应的hash函数，比如默认情况下，Key为string时就有，Key为int时就没有。

　　对于我们只增不删的特殊场景，_Hash_node创建后就一直存活，而指针数组_M_buckets会在rehash时重新开辟更大的指针数组空间，然后把每个_Hash_node挂在新的桶上构成新的链表，最后释放旧的指针数组空间。每当元素数量达到桶的数量时就会触发rehash，桶的数量按照11、23、47、97、199、409、823、1741、3739、7517、15173这样大约两倍的规模扩张，且桶数一定是质数。

　　综上，unordered_map中的Alloc既要管_Hash_node的分配，也要管_M_buckets的分配和释放，而不是像表面看起来负责pair<const Key,T>的分配。

　　unordered_map默认的Alloc是std::allocator，内部通过rebind技巧可以把allocator<pair<const Key,T> >类型的分配器重绑定出allocator<_Hash_node>和allocator<_Hash_node_base*>的分配器。

　　std::allocator继承自__gnu_cxx::new_allocator，包含两个成员函数allocate和deallocate，顾名思义可知一个负责分配一个负责释放，实现方法就是最基础的::operator new和::operator delete，因此在大量申请_Hash_node内存时一样会出现之前说过的问题：glibc申请的内存比预想的要大得多。我们只好接管unordered_map的Alloc，实现自定义的my_allocator，同样在mem_pool上分配_Hash_node的内存。

　　这里有个问题，我们只想接管_Hash_node的分配，而_M_buckets的分配依然使用std::allocator（因为M_buckets的内存分配每次rehash会有释放过程，mem_pool不再适用，好在次数不多，且后面分配的内存越来越大，用::operator new也没太大问题 ），但unordered_map的定义限制了只能传入一种分配器，那就只好在函数allocate上做文章，通过my_allocator模板参数T的类型来判断当前到底是给谁分配，然后区别对待，代码就像下面这样：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
template<typename T>
class my_allocator : public allocator<T>
{
    ... // other lines
    pointer allocate(size_type count)
    {
        if(typeid(T) == typeid(__detail::_Hash_node_base*))
        {
            return allocator<T>::allocate(count);
        }
        return (pointer)mem_pool::get_mem(sizeof(T) * count);
    }

    void deallocate(pointer ptr, size_type count)
    {
        if(typeid(T) == typeid(__detail::_Hash_node_base*))
        {
            allocator<T>::deallocate(ptr, count);
        }
    }
};

6. 去除_Hash_node中的_M_hash_code

　　第5点提到了struct _Hash_node中可能会包含一项_M_hash_code，目的是缓存hash值提高计算性能。当Key为string类型时是包含这一项的，第4点提到我们把string改成了char*，同时实现自定义的仿函数my_hash，结果发现也会包含这一项。这当然是不能忍的，违背了我们一切以节约内存为先的最高原则。我各种实验，先是如此冗余地解决了：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
namespace std
{
    template<>
    struct __is_fast_hash<my_hash> : public std::true_type
    {};
    template<>
    struct is_default_constructible<my_hash> : public std::true_type
    {};
    template<>
    struct is_copy_assignable<my_hash> : public std::true_type
    {};
namespace __detail
{
    template<>
    struct __is_noexcept_hash<const char*, my_hash> : public std::true_type
    {};
}
}

　　后来发现只要在my_hash的operator()里加上C++11关键字noexcept即可：

1
2
3
4
5
6
7
8
class my_hash
{
public:
    size_t operator()(const char* const& key) const noexcept
    {
        return _Hash_impl::hash(key, strlen(key));
    }
};

7. unordered_map的value类型从ftrl_model_unit*换成ftrl_model_unit

　　之前因为ftrl_model_unit的分配已被接管到mem_pool，而unordered_map还是缺省归std::allocator管，所以unordered_map里的value存的是指针ftrl_model_unit*，现在既然unordered_map的内存分配也被我们包办了，这个指针也就下岗了，直接把ftrl_model_unit放进unordered_map里，又能省个指针的空间，8字节呢。现在unordered_map变成：

1
using my_hash_map = unordered_map<const char*, ftrl_model_unit, my_hash, my_equal, my_allocator< pair<const char*, ftrl_model_unit> > >;

　　结点_Hash_node的内存布局像这样：

1
2
3
4
5
6
struct _Hash_node
{
    _Hash_node_base* _M_nxt;
    const char* Key;
    ftrl_model_unit Value;
};

　　需要修改my_allocator的allocate函数：

1
2
3
4
5
6
7
8
9
pointer allocate(size_type count)
{
    if(typeid(T) == typeid(__detail::_Hash_node_base*))
    {
        return allocator<T>::allocate(count);
    }
    // here, count is 1, T is std::__detail::_Hash_node<std::pair<const char* const, ftrl_model_unit>, false>
    return (pointer)mem_pool::get_mem(sizeof(T) + ftrl_model_unit::get_ext_mem_size());
}

　　申请的内存除了_Hash_node的大小，还要额外的内存放ftrl_model_unit对应的vi，v_ni，v_zi三个数组，目的见第3条。

8. ftrl_model_unit改成模板类，支持选择double或float

　　为了进一步节省，模型参数可以选择用float存储。在训练参数里加了-mnt选项，默认为double，可以指定为float。具体实现就是代码里大部分class都改成了模板类，比如：

1
2
3
4
5
6
7
8
9
template<typename T>
class ftrl_model_unit
{
public:
    T wi;
    T w_ni;
    T w_zi;
    ...// other lines
};

　　毕竟sizeof(float) = 4而sizeof(double) = 8，能省一半当然好，但float范围小，且精度只有6~7位，而double精度高达15~16位，因此要慎重使用，可能会影响模型效果。

9. _Hash_node中padding的处理

　　ftrl_model_unit的模板参数是double时，_Hash_node的内存布局如下：

1
2
3
4
5
6
struct _Hash_node                   // 40 bytes
{
    _Hash_node_base* _M_nxt;        // 8 bytes
    const char* Key;                // 8 bytes
    ftrl_model_unit<double> Value;  // 8*3 bytes = 24 bytes
};

　　可见_Hash_node很“紧实”，没有空隙。当改成float后：

1
2
3
4
5
6
struct _Hash_node                   // 32 bytes
{
    _Hash_node_base* _M_nxt;        // 8 bytes
    const char* Key;                // 8 bytes
    ftrl_model_unit<float> Value;   // 4*3 bytes = 12 bytes
};

　　出现了4字节的空隙，相当于这样：

1
2
3
4
5
6
7
struct _Hash_node                   // 32 bytes
{
    _Hash_node_base* _M_nxt;        // 8 bytes
    const char* Key;                // 8 bytes
    ftrl_model_unit<float> Value;   // 4*3 bytes = 12 bytes
    char padding[4];                // 4 bytes
};

　　这是由于64位上默认对齐系数为8导致的，_Hash_node需要填充4字节凑成8的倍数。

　　如果vi，v_ni，v_zi三个数组的额外空间起始位置是跟在_Hash_node后面，那么这4字节就浪费掉了。特别是当factor_num=0，即没有这三个数组时，这4字节也是浪费。本着寸土必争的精神，必须把这4字节的内存省下来。一开始想到的是用#pragma pack (1)把struct _Hash_node压紧实：

1
2
3
#pragma pack (1)
#include <unordered_map>
#pragma pack ()

　　在独立的测试代码上发现这么做可以达到目的，但在alphaFM上没这么蛮干，毕竟#include <unordered_map>里会引入不少头文件，直接套个#pragma pack (1)会改变很多的STL类内存布局，总担心会带来预想不到的隐患。

　　替代方案依然是细致地控制内存分配和布局，比如float版本的_Hash_node大小是32字节，假设factor_num=2，那么三个数组总共大小是3*factor_num*sizeof(float)=3*2*4=24。原来的allocate函数会申请32+24=56字节，其中0~31共32字节放_Hash_node，32~55共24字节放三个数组；新的方案只会申请52字节，0~31依然放_Hash_node，但三个数组从offset 28开始，即28~51放三个数组，和_Hash_node有重叠，正好利用4字节的padding部分。特别地，当factor_num=0时，allocate就只申请28字节内存。

　　至此，fm_train的内存优化之路就走到这里，内存消耗不再是一笔糊涂账，在任务启动前就可以大致预估。设特征维度为d，特征名字符串平常长度（包括结尾\0）为s，模型参数类型为T，则内存消耗的大户包括：

　　(1) mem_pool上分配特征名字符串，共s*d

　　(2) mem_pool上分配_Hash_node和vi、v_ni、v_zi三个数组，共(8+8+3*sizeof(T)+3*factor_num*sizeof(T))*d

　　(3) std::allocator分配的_M_buckets，考虑到rehash时候的内存消耗，共8*1.5*d/load_factor，目前没有修改max_load_factor，还是默认值1，所以load_factor大约分布在0.5~1之间，可得这一项消耗为12*d到24*d之间。如果增大max_load_factor，这里还可以优化，为了效率暂时没动。

　　加起来一共是(28+3*(1+factor_num)*sizeof(T)+s)*d到(40+3*(1+factor_num)*sizeof(T)+s)*d

　　比如d为10亿即大概是1G，s=35，factor_num=0，T为double，则内存消耗为(28+3*8+35)G=87G到(40+3*8+35)G=99G，再加上一些额外的消耗，最多应该100G左右。

10. 优化fm_predict的内存占用

　　优化fm_train之后，再来优化fm_predict。之前偷懒，预测时也会把完整模型加载到内存，导致内存消耗基本和fm_train一致，其实对于预测只需要加载非零的wi和vi项，其他的w_ni、w_zi、v_ni、v_zi都不需要，这样优化后，内存消耗和fm_train比基本可以忽略了。

11. gcc的版本兼容性问题

　　代码中涉及到了gcc的STL具体实现，而gcc不同的版本之间实现代码还有差异，真是一个糟心的问题。我对比测试了4.8.5和5.4.0以及7.3.0版本，为了兼容它们，利用了C++模板一种称作SFINAE的“奇技淫巧”，具体不展开了。之前看到SFINAE的时候觉得这玩意儿完全无实用价值，没想到这次就用到了，啪啪打脸。

　　gcc版本太多，我不可能每个版本都编译测试一下，只能乐观假定4.8.5和7.3.0之间的都没问题，好吧，这算是线性插值的思路？

12. one more thing，模型文件增加二进制格式

　　在实践中当模型维度很高时，一个痛点是耗内存，另一个痛点是模型加载和输出的时间特别慢，有时能到几十分钟。

　　之前模型文件只有文本格式，所以这次一并做了优化，加入了二进制格式的选项，加载和输出时间大大加快，能有10倍量级的加速。

　　同时提供了一个模型文件的格式转换工具model_bin_tool，可以查看模型信息，可以在二进制和文本格式之间互转，方便模型文件的后续上线使用。model_bin_tool的具体用法参见README即可。