Float16 TFLite models with multiple interpreters doesn't share weights but bloats memory

[barabanus]

System information

• Linux 5.11.4-arch1-1 Add support for Python 3.x #1 SMP PREEMPT Sun, 07 Mar 2021 18:00:49 +0000 x86_64 GNU/Linux
• TensorFlow 2.4.1 from tensorflow-opt-cuda package (https://archlinux.org/packages/community/x86_64/tensorflow-opt-cuda)

Describe the current behavior

When we create multiple interpreters per single model (e.g. for thread-safety reasons) the expected behavior is that all interpreters share model weights. And that's statement is true when the model is float32 model:

$ ./minimal mem-bloat-float32.tflite        
initial mem 1 MB
interpreter #1, mem 139 MB
interpreter #2, mem 141 MB
interpreter #3, mem 144 MB
interpreter #4, mem 146 MB
interpreter #5, mem 148 MB
interpreter #6, mem 150 MB
interpreter #7, mem 152 MB
interpreter #8, mem 154 MB
interpreter #9, mem 156 MB

But when the same model had been converted to be float16 the weights are not shared anymore and memory is bloated so that each interpreter has its own copy of model weights:

$ ./minimal mem-bloat-float16.tflite
initial mem 1 MB
interpreter #1, mem 206 MB
interpreter #2, mem 341 MB
interpreter #3, mem 476 MB
interpreter #4, mem 610 MB
interpreter #5, mem 745 MB
interpreter #6, mem 880 MB
interpreter #7, mem 1015 MB
interpreter #8, mem 1149 MB
interpreter #9, mem 1284 MB

Describe the expected behavior

The expected behavior is that multiple interpreters share weights for the same model, whenever it's float16 or float32 model.

Standalone code to reproduce the issue

• unzip TFLite models: mem-bloat-tflite.zip
• build TFLite minimal example with the following code within tensorflow/tensorflow/lite/examples/minimal.cc:

#include <cstdio>

#include "tensorflow/lite/interpreter.h"
#include "tensorflow/lite/kernels/register.h"
#include "tensorflow/lite/model.h"
#include "tensorflow/lite/optional_debug_tools.h"

#if defined(_WIN32)
#include <windows.h>
#include <psapi.h>

#elif defined(__unix__) || defined(__unix) || defined(unix) || (defined(__APPLE__) && defined(__MACH__))
#include <unistd.h>
#include <sys/resource.h>

#if defined(__APPLE__) && defined(__MACH__)
#include <mach/mach.h>

#elif (defined(_AIX) || defined(__TOS__AIX__)) || (defined(__sun__) || defined(__sun) || defined(sun) && (defined(__SVR4) || defined(__svr4__)))
#include <fcntl.h>
#include <procfs.h>

#elif defined(__linux__) || defined(__linux) || defined(linux) || defined(__gnu_linux__)
#include <stdio.h>

#endif

#else
#error "Cannot define getCurrentRSS( ) for an unknown OS."
#endif


/**
 * Returns the current resident set size (physical memory use) measured
 * in bytes, or zero if the value cannot be determined on this OS.
 */
size_t getCurrentRSS( )
{
#if defined(_WIN32)
    /* Windows -------------------------------------------------- */
    PROCESS_MEMORY_COUNTERS info;
    GetProcessMemoryInfo( GetCurrentProcess( ), &info, sizeof(info) );
    return (size_t)info.WorkingSetSize;

#elif defined(__APPLE__) && defined(__MACH__)
    /* OSX ------------------------------------------------------ */
    struct mach_task_basic_info info;
    mach_msg_type_number_t infoCount = MACH_TASK_BASIC_INFO_COUNT;
    if ( task_info( mach_task_self( ), MACH_TASK_BASIC_INFO,
        (task_info_t)&info, &infoCount ) != KERN_SUCCESS )
        return (size_t)0L;      /* Can't access? */
    return (size_t)info.resident_size;

#elif defined(__linux__) || defined(__linux) || defined(linux) || defined(__gnu_linux__)
    /* Linux ---------------------------------------------------- */
    long rss = 0L;
    FILE* fp = NULL;
    if ( (fp = fopen( "/proc/self/statm", "r" )) == NULL )
        return (size_t)0L;      /* Can't open? */
    if ( fscanf( fp, "%*s%ld", &rss ) != 1 )
    {
        fclose( fp );
        return (size_t)0L;      /* Can't read? */
    }
    fclose( fp );
    return (size_t)rss * (size_t)sysconf( _SC_PAGESIZE);

#else
    /* AIX, BSD, Solaris, and Unknown OS ------------------------ */
    return (size_t)0L;          /* Unsupported. */
#endif
}


#define TFLITE_MINIMAL_CHECK(x)                              \
  if (!(x)) {                                                \
    fprintf(stderr, "Error at %s:%d\n", __FILE__, __LINE__); \
    exit(1);                                                 \
  }

int main(int argc, char* argv[])
{
  if (argc != 2) {
    fprintf(stderr, "minimal <tflite model>\n");
    return 1;
  }
  const char* filename = argv[1];

  printf("initial mem %d MB\n", getCurrentRSS() >> 20);

  // Load model
  std::unique_ptr<tflite::FlatBufferModel> model = tflite::FlatBufferModel::BuildFromFile(filename);
  TFLITE_MINIMAL_CHECK(model != nullptr);

  tflite::ops::builtin::BuiltinOpResolver resolver;
  tflite::InterpreterBuilder builder(*model, resolver);

  std::unique_ptr<tflite::Interpreter> interpreter_list[9];
  for (auto &interpreter : interpreter_list) {

    builder(&interpreter, 1);
    TFLITE_MINIMAL_CHECK(interpreter != nullptr);
    TFLITE_MINIMAL_CHECK(interpreter->AllocateTensors() == kTfLiteOk);
    interpreter->Invoke();

    printf("interpreter #%d, mem %d MB\n", &interpreter - &interpreter_list[0] + 1, getCurrentRSS() >> 20);
  }

  return 0;
}

Other info / logs
Both versions of model had been converted from this protobuf file: mem-bloat-protobuf.zip

[abattery]

Thanks for your valuable feedback. Currently, each TensorFlow Lite interpreter owns each memory space by design. We will consider this as a feature request for increasing memory sharing across the multiple interpreter interfaces.

@daverim @terryheo any thoughts on this?

[terryheo]

Interpreter can share data in flatbuffer file which contains constant tensors.
But each interpreter needs its own memory to store intermediate tensors. so it's natural to see a certain amount of memory increasement.

Regarding your testing, I have few comments.

1. Could you check /proc/self/statm shared pages?
   According to https://www.kernel.org/doc/html/latest/filesystems/proc.html, VmRSS = RssAnon + RssFile + RssShmem
   which means VmRSS contains RssFile size. So you need to check it separately.
2. If RssFile is zero, please confirm your model is using MMAPAllocation() https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/model_builder.cc#L50

[barabanus]

@terryheo

Interpreter can share data in flatbuffer file which contains constant tensors. But each interpreter needs its own memory to store intermediate tensors. so it's natural to see a certain amount of memory increasement.

That's right, but the difference between models is datatype (float32 vs float16), so it's counter-intuitive to see that float32 version do share weights and float16 does not.

The memory usage gives insights of what's going on with memory usage.

For float32 we've got: 1, 139, 141, 144, 146 etc.
That means 137 MB of shared constant tensors in float32 + 2 MB for intermediate tensors per interpreter.
So model of the memory usage is (roughly): 137 + 2 * interpreters_count MB

For float16 we've got: 1, 206, 341, 476, 610 etc.
That means 69 MB of constant tensors in float16 + 137 MB for intermediate tensors per interpreter
So model of the memory usage is (roughly): 69 + 137 * interpreters_count MB
..and 137 MB for intermediate tensors looks very similar to 137 MB of the float32 model, so I assume there's conversion of float16 constant tensors to float32 constant tensors within each interpreter.

Could you check /proc/self/statm shared pages?
please confirm your model is using MMAPAllocation()

I modified both the program and TFLite sources to show /proc/self/statm and MMAPAllocation::IsSupported() (at the location you pointed)

$ ./minimal mem-bloat-float32.tflite                        
initial mem 1 MB              [/proc/self/statm -> 2289 486 432 726 0 105 0]
MMAPAllocation::IsSupported() is false
interpreter #1, mem  140 MB   [/proc/self/statm -> 36798 35848 1341 726 0 34614 0]
interpreter #2, mem  142 MB   [/proc/self/statm -> 37329 36442 1341 726 0 35145 0]
interpreter #3, mem  144 MB   [/proc/self/statm -> 37893 36970 1341 726 0 35709 0]
interpreter #4, mem  146 MB   [/proc/self/statm -> 38424 37498 1341 726 0 36240 0]
interpreter #5, mem  148 MB   [/proc/self/statm -> 38988 38026 1341 726 0 36804 0]
interpreter #6, mem  150 MB   [/proc/self/statm -> 39519 38554 1341 726 0 37335 0]
interpreter #7, mem  152 MB   [/proc/self/statm -> 40083 39148 1341 726 0 37899 0]
interpreter #8, mem  154 MB   [/proc/self/statm -> 40614 39676 1341 726 0 38430 0]
interpreter #9, mem  157 MB   [/proc/self/statm -> 41178 40204 1341 726 0 38994 0]
$ ./minimal mem-bloat-float16.tflite                        
initial mem 1 MB              [/proc/self/statm -> 2289 501 447 726 0 105 0]
MMAPAllocation::IsSupported() is false
interpreter #1, mem  206 MB   [/proc/self/statm -> 53770 52848 1301 726 0 51586 0]
interpreter #2, mem  341 MB   [/proc/self/statm -> 88268 87352 1301 726 0 86084 0]
interpreter #3, mem  476 MB   [/proc/self/statm -> 122811 121863 1301 726 0 120627 0]
interpreter #4, mem  610 MB   [/proc/self/statm -> 157284 156354 1301 726 0 155100 0]
interpreter #5, mem  745 MB   [/proc/self/statm -> 191790 190842 1301 726 0 189606 0]
interpreter #6, mem  880 MB   [/proc/self/statm -> 226263 225332 1301 726 0 224079 0]
interpreter #7, mem 1014 MB   [/proc/self/statm -> 260736 259820 1301 726 0 258552 0]
interpreter #8, mem 1149 MB   [/proc/self/statm -> 295202 294286 1301 726 0 293018 0]
interpreter #9, mem 1284 MB   [/proc/self/statm -> 329746 328798 1301 726 0 327562 0]

[barabanus]

@terryheo
I rebuilt the example with -DTFLITE_ENABLE_MMAP=ON and now MMAPAllocation::IsSupported() is true, but memory bloat for float16 is still here:

$ ./minimal mem-bloat-float32.tflite                   
initial mem 1 MB              [/proc/self/statm -> 2291 486 432 727 0 105 0]
MMAPAllocation::IsSupported() is true
interpreter #1, mem   10 MB   [/proc/self/statm -> 36800 2583 1974 727 0 670 0]
interpreter #2, mem   12 MB   [/proc/self/statm -> 37331 3111 1974 727 0 1201 0]
interpreter #3, mem   14 MB   [/proc/self/statm -> 37895 3639 1974 727 0 1765 0]
interpreter #4, mem   16 MB   [/proc/self/statm -> 38426 4167 1974 727 0 2296 0]
interpreter #5, mem   18 MB   [/proc/self/statm -> 38990 4761 1974 727 0 2860 0]
interpreter #6, mem   20 MB   [/proc/self/statm -> 39521 5289 1974 727 0 3391 0]
interpreter #7, mem   22 MB   [/proc/self/statm -> 40085 5817 1974 727 0 3955 0]
interpreter #8, mem   24 MB   [/proc/self/statm -> 40616 6345 1974 727 0 4486 0]
interpreter #9, mem   26 MB   [/proc/self/statm -> 41180 6873 1974 727 0 5050 0]
$ ./minimal mem-bloat-float16.tflite
initial mem 1 MB              [/proc/self/statm -> 2291 501 447 727 0 105 0]
MMAPAllocation::IsSupported() is true
interpreter #1, mem  206 MB   [/proc/self/statm -> 53772 52834 18263 727 0 34612 0]
interpreter #2, mem  341 MB   [/proc/self/statm -> 88269 87354 18279 727 0 69109 0]
interpreter #3, mem  476 MB   [/proc/self/statm -> 122813 121866 18279 727 0 103653 0]
interpreter #4, mem  610 MB   [/proc/self/statm -> 157286 156357 18279 727 0 138126 0]
interpreter #5, mem  745 MB   [/proc/self/statm -> 191792 190845 18279 727 0 172632 0]
interpreter #6, mem  880 MB   [/proc/self/statm -> 226265 225335 18279 727 0 207105 0]
interpreter #7, mem 1014 MB   [/proc/self/statm -> 260738 259823 18279 727 0 241578 0]
interpreter #8, mem 1149 MB   [/proc/self/statm -> 295204 294289 18279 727 0 276044 0]
interpreter #9, mem 1284 MB   [/proc/self/statm -> 329748 328801 18279 727 0 310588 0]

[barabanus]

Let me humbly emphasize that weights sharing between interpreters is very important. With our runtime environment we're running up to eight threads of realtime inferences with precomputed ~130 MB of float32 weights. It's counter-productive to copy ~130 MB of weights per each interpreter. And also it would be better to have ~65 MB of float16 weights shared between eight interpreters/threads.

[terryheo]

..and 137 MB for intermediate tensors looks very similar to 137 MB of the float32 model, so I assume there's conversion of float16 constant tensors to float32 constant tensors within each interpreter.

I think this is the case since there is no float16 native kernels. @jdduke, @daverim FYI,

I wonder if you tried full integer quantization which saves more memory. https://www.tensorflow.org/lite/performance/post_training_quantization
I guess it also allows to share constant tensors.

[barabanus]

@terryheo

I wonder if you tried full integer quantization which saves more memory

At some point I will, but I'd split the whole graph into preparation (float32), inference (int quantization) and postprocessing (float32) parts. Currently it would be very beneficial to have just float16. More to say, it's shame that previous implementation of the same model had been done manually with float16 arithmetic (== two times less memory) and the same performance. Also, for a single thread there was no need to have 2 MB of intermediate tensors per interpreter. So I hope TFLite will improve in this direction.

Although TFLite has no native float16 kernels why not introduce shared memory manager to each model instance? When the first interpreter converts float16 into float32 constant tensors they are registered within the shared memory manager. When the next interpreter requires float32 constants it gets its smart pointer from the manager. (For example, by unique name of tensor or node ID).

[jdduke]

Hi @barabanus . Your request is quite reasonable, though there would be quite a bit of complexity involved to make this work properly, due to the lazy manner in which float16 weights are currently "inflated" at runtime. Terry's suggestion to try int8 quantization is probably the best approach to both reduce the model size but also enable weight sharing.

That said, for Q2 we are looking into supporting native float16 execution paths that operate natively on float16 weights for ARM devices. There are some issues involved as it relates to accuracy (the most widely available float16 ARM intrinsics do not accumulate into float32, though the latest devices to support this), but we'll also be working on tooling to validate that the relaxed runtime precision doesn't significantly degrade accuracy.