Implementation:Microsoft DeepSpeedExamples Gather Memory

Knowledge Sources	Microsoft_DeepSpeedExamples
Domains	Memory Profiling, Log Parsing
Last Updated	2026-02-07 12:00 GMT

Overview

A command-line script that parses training log files to gather, compare, and summarize GPU memory usage across different BF16 master weight configurations.

Description

gather_memory.py is a log-parsing utility designed for analyzing memory consumption in BF16 low-precision master weight training experiments. The script reads structured log files from a specified directory, extracts SUMMARY lines using regular expressions, and compiles a comparison table showing peak memory, allocated memory, average step time, and memory reduction percentages relative to a baseline configuration.

The script supports two primary log files: baseline.log and bf16_full.log. It parses each file for a standardized summary line format containing config name, parameter count, peak memory bytes, allocated memory bytes, and average step time. The results are then formatted into both a human-readable text table and a Markdown table suitable for inclusion in README documentation.

Output is printed to stdout and optionally saved to a file (defaulting to summary.txt in the log directory). The script provides detailed per-configuration breakdowns including formatted byte values in GB and MB, as well as percentage memory reduction calculations compared to the baseline fp32 master weight configuration.

Usage

Use this script after running BF16 master weight training experiments to compare memory usage between baseline (fp32 master weights) and BF16 low-precision configurations. It is particularly useful for generating benchmark tables for documentation and for quick analysis of memory savings achieved through lower-precision master weights.

Code Reference

Source Location

• Repository: Microsoft_DeepSpeedExamples
• File: training/bf16_master_weight/gather_memory.py
• Lines: 1-175

Signature

def parse_summary_line(line):
    """Parse the SUMMARY line from log output."""

def format_bytes(bytes_val):
    """Format bytes to human-readable string."""

def format_bytes_mb(bytes_val):
    """Format bytes to MB."""

def get_config_name(config_path):
    """Extract clean config name from path."""

def main():
    parser = argparse.ArgumentParser(description="Gather memory usage from training logs")
    parser.add_argument("--log_dir", type=str, required=True, help="Directory containing log files")
    parser.add_argument("--output", type=str, default=None, help="Output file for summary")

Import

from training.bf16_master_weight.gather_memory import parse_summary_line, format_bytes, main

I/O Contract

Inputs

Name	Type	Required	Description
--log_dir	str	Yes	Directory containing training log files (baseline.log, bf16_full.log)
--output	str	No	Output file path for summary; defaults to log_dir/summary.txt

Outputs

Name	Type	Description
summary text	str (stdout)	Formatted comparison table printed to console
summary file	file	Text file containing the full memory comparison report with both plain-text and Markdown tables
return code	int	0 on success, 1 on error (missing directory or no results)

Usage Examples

# Run from command line to compare memory usage
# python gather_memory.py --log_dir logs/20231201_120000

# Specify custom output location
# python gather_memory.py --log_dir logs/20231201_120000 --output results/memory_comparison.txt

# Use parse_summary_line programmatically
from training.bf16_master_weight.gather_memory import parse_summary_line

line = "SUMMARY: config=baseline params=110000000 peak_mem_bytes=4294967296 alloc_mem_bytes=2147483648 avg_step_time=0.5432"
result = parse_summary_line(line)
# result = {
#     "config": "baseline",
#     "params": 110000000,
#     "peak_mem_bytes": 4294967296,
#     "alloc_mem_bytes": 2147483648,
#     "avg_step_time": 0.5432
# }