Options for Stagger model loading for low memory systems

[jjhursey]

• --stagger_load : (default: 0 off) Stagger model loading to avoid OOM issues on the host
• --stagger_update_lazyhandle : (default: 0 off) Stagger update_lazyhandle to avoid OOM issues on the host
• --dist_timeout : (default: either 10 for NCCL or 30 for others set by PyTorch) torch distributed timeout in minutes

[jjhursey]

Keeping this in Draft for now while a few teams get back to me on testing.

[thanh-lam]

I checked out this PR and am testing it with latest FMS levels. Reasons for trying this:

• We have been hitting failures with "Resource temporarily unavailable" when running granite-3.0-8b with 8-AIUs. This is an example of a large model but there're other models hitting this failures as well.
• With this PR, I added the two new parameters when invoking inference.py:
  • --stagger_load 1
  • --stagger_update_lazyhandle 1
• So far, tests are passing without hitting "Resource temporarily unavailable"

[thanh-lam]

For testing purposes: What are differences between option 2 and 1?
@jjhursey Can you add some descriptions for these options?

[jjhursey]

For testing purposes: What are differences between option 2 and 1?

We can improve the help text, but this is what it reports now:

    "--stagger_load",
    help="Stagger model loading to avoid OOM issues on the host"

    "--stagger_update_lazyhandle",
    help="Stagger update_lazyhandle to avoid OOM issues on the host"

So two different options to stagger processing of two different sections of the compile phase in case you have a need to set different values for each part.

The default value is 0 which means skip the staggering and let all of the processes enter this section at the same time. This results in the maximum amount of memory utilization, but also (generally) the fastest time through the compile phase.

Setting this value to >0 defines how many concurrent processes are allowed to be in that phase at the same time. So if you run with 16 processes and set the value to 2 then only 2 processes at a time will be performing the compile phase. Once the first set of 2 are done then the next set of 2 will proceed and so on. All processes that have completed the compile will wait for the others to finish before moving on to the next section of the code.

Setting the value to 1 means only one process at a time. This will use the least amount of memory, but will result in the longest compile time since we are only compiling for one process at a time. Increasing the value increases the total amount of memory, but should decrease the compile time.

If you are only worried about functional testing in a memory constrained environment and not on compile time efficiency then setting this value to 1 is likely what you want.

[thanh-lam]

@jjhursey , thanks! This is very helpful.
Any estimate when or will these new parameters be added to inference.py?
Because this "helps" get around the model failure: "Resource temporarily unavailable" with 8-AIU tests, we need to document it while waiting for "real" memory fixes.

[jjhursey]

Both the x86 and Power test teams have confirmed that this helps mitigate their testing needs for large models on low memory systems.

[jjhursey]

@JRosenkranz @ani300 this is ready for review

[JRosenkranz]

It looks like logic is called multiple times, do you think it would make sense to put it in it's own utility function, this way in can be re-used in the future

[jjhursey]

Yeah I can do that. Have an "enter" and "exit" version to place in the code.

[jjhursey]

I just pushed a commit for this change

[JRosenkranz]

Has this been tested with inference.py as well as test_decoders (multi-aiu / single aiu). In theory those should not change. Also, it might make sense to add an option to this in test_decoders in low memory systems.

[JRosenkranz]

bot:test
TEST_FILE=test_decoders.py MODEL_ID=ibm-granite/granite-3.2-8b-instruct BATCH_SIZE=8 SEQUENCE_LENGTH=64 USE_TINY_MODEL=0 NUM_AIU=2

[dpatel-ops]

bot:test
TEST_FILE=test_decoders.py MODEL_ID=ibm-granite/granite-3.2-8b-instruct BATCH_SIZE=8 SEQUENCE_LENGTH=64 USE_TINY_MODEL=0 NUM_AIU=2

[jjhursey]

I pushed a commit to consolidate the staggered enter/exit. I also rebased on main

[JRosenkranz]

bot:test
TEST_FILE=test_decoders.py MODEL_ID=ibm-granite/granite-3.2-8b-instruct BATCH_SIZE=8 SEQUENCE_LENGTH=64 USE_TINY_MODEL=0

[JRosenkranz]

bot:test
TEST_FILE=test_decoders.py MODEL_ID=ibm-granite/granite-3.2-8b-instruct BATCH_SIZE=8 SEQUENCE_LENGTH=64 USE_TINY_MODEL=0

[JRosenkranz]

bot:test
TEST_FILE=test_decoders.py MODEL_ID=ibm-granite/granite-3.2-8b-instruct BATCH_SIZE=8 SEQUENCE_LENGTH=64 USE_TINY_MODEL=0

[jjhursey]

I pushed an update that adds the docstrings and fixes a DCO check.

[JRosenkranz]

Would it make sense for this to be a stagger_context (given each stagger_enter needs to be paired with a stagger_leave):

with stagger_context(limit):
     model = get_model(...)

[jjhursey]

Yeah, that would be a good idea. I've not built a context like that in Python before, but I'm always willing to learn. I'll take a pass at it in the next couple of days.

[jjhursey]

I added a commit to convert this to a contextlib function. Take a look and let me know what you think.

[jjhursey]

I recently updated my foundation-model-stack repo and the rest of the stack, and now I'm seeing a hang during the generate wrapped by with torch_sendnn.warmup_mode(): because the second iteration is calling allreduce. So the allreduce is hanging because not all of the processes are participating in the warmup at the same time.

I'm not sure what part of the stack is causing that to break. It's not caused by this PR, but a new synchronization in the stack it is enclosing in the warmup.