PyTorch memory profiler memory timeline not showing categories

drwslacy47: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/a/ea69e8b74a24c575b6bd128631ef2eebf6042dc8.png" data-download-href="https://discuss.pytorch.org/uploads/default/ea69e8b74a24c575b6bd128631ef2eebf6042dc8" title="no_categories"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/a/ea69e8b74a24c575b6bd128631ef2eebf6042dc8_2_690x429.png" alt="no_categories" data-base62-sha1="xrIDvhLj9eb537nesmMK9OWGsxa" width="690" height="429" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/a/ea69e8b74a24c575b6bd128631ef2eebf6042dc8_2_690x429.png, https://discuss.pytorch.org/uploads/default/optimized/3X/e/a/ea69e8b74a24c575b6bd128631ef2eebf6042dc8_2_1035x643.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/e/a/ea69e8b74a24c575b6bd128631ef2eebf6042dc8_2_1380x858.png 2x" data-dominant-color="F3F3F3"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">no_categories</span><span class="informations">1493×929 81.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I’m trying to use PyTorch’s memory timeline generated by the profiler to visualize what is contributing to a GPU OOM problem.  The memory allocation ramp shown in the attached image is happening during the first forward pass of a 13B parameter Llama2 model.  I don’t understand why the memory allocations are categorized as “unknown” instead of using the other categories shown in the legend (e.g., parameters, activations, etc).</p>
<p>One piece of info that may or may not be relevant:  I’m using DeepSpeed Zero-3 on a 16-GPU cluster.  Is it possible that DeepSpeed isn’t playing well with the profiler’s ability to recognize the memory allocation categories?</p>
<p>Here’s my code…</p>
<pre><code class="lang-auto">def TrainSampleBasedDataParallel(gargs, clargs):
    # Create and initialize model
    model = CreateAndInitializeLlama2Model(gargs)
    DS.runtime.zero.stage3.estimate_zero3_model_states_mem_needs_all_live(model, num_gpus_per_node=16, num_nodes=1)

    modelLayers = model.AsPipelineLayers()

    modelWithObjectives = ModelWithObjectives(modelLayers)
    DS.zero.Init(module=modelWithObjectives)

    modelEngine, _, _, _ = DS.initialize(config=gargs.DEEPSPEED_CONFIG,
                                            model=modelWithObjectives,
                                            model_parameters=modelWithObjectives.parameters()
                                        )


    # Create Dataloaders
    trainDataLoader, evalDataLoader = CreateDataLoaders(gargs)

    import logging
    import socket
    from datetime import datetime, timedelta

    from torch.autograd.profiler import record_function
    from torchvision import models

    logging.basicConfig(
    format="%(levelname)s:%(asctime)s %(message)s",
    level=logging.INFO,
    datefmt="%Y-%m-%d %H:%M:%S",
    )
    logger: logging.Logger = logging.getLogger(__name__)
    logger.setLevel(level=logging.INFO)

    TIME_FORMAT_STR: str = "%b_%d_%H_%M_%S"

    def trace_handler(prof: PT.profiler.profile):
        # Prefix for file names.
        host_name = socket.gethostname()
        timestamp = datetime.now().strftime(TIME_FORMAT_STR)
        file_prefix = f"{host_name}_{timestamp}_gpu{gargs.DEVICE_ID}"

        # Construct the trace file.
        prof.export_chrome_trace(f"{file_prefix}.json.gz")

        # Construct the memory timeline file.
        prof.export_memory_timeline(f"{file_prefix}.html", device=gargs.DEVICE)

    with PT.profiler.profile(
        activities=[
            PT.profiler.ProfilerActivity.CPU,
            PT.profiler.ProfilerActivity.CUDA,
        ],
        schedule=PT.profiler.schedule(wait=0, warmup=0, active=6, repeat=1),
        record_shapes=True,
        profile_memory=True,
        with_stack=True,
        on_trace_ready=trace_handler,
    ) as prof:
        # ================
        # Training loop
        # ================    
        modelEngine.zero_grad()
        for trainBatch in trainDataLoader:
            prof.step()
            with record_function("## forward ##"):
                trainBatchOnDevice = SendTupleOfTupleTensorsToDevice(trainBatch, gargs.DEVICE)
                trainLoss, output = modelEngine(trainBatchOnDevice)
            with record_function("## backward ##"):
                modelEngine.backward(trainLoss)
            with record_function("## optimizer ##"):
                modelEngine.step()
                modelEngine.zero_grad()
</code></pre>