Out of Memory Issue With Segmented Output

[Z-Richard]

Example inference_config.yaml:

experiment_dir: /workspace_hdd/haoyu/ACE2-inference-parent/ACE2-inference-2020_2025_init_feb01_segmented_output
n_forward_steps: 146
forward_steps_in_memory: 50
checkpoint_path: /workspace_hdd/haoyu/ACE2-ERA5/ace2_era5_ckpt.tar
logging:
  log_to_screen: true
  log_to_wandb: false
  log_to_file: true
  project: ace
initial_condition:
  path: /workspace_hdd/haoyu/ACE2-ERA5/initial_conditions/ic_2020.nc
  start_indices:
    times:
      - "2020-02-01T00:00:00"
forcing_loader:
  dataset:
    data_path: /workspace_hdd/haoyu/ACE2-ERA5/forcing_data
  num_data_workers: 4
data_writer:
  save_prediction_files: true
  save_monthly_files: false
  names: ["PRATEsfc"]

Command: nohup python -m fme.ace.inference inference_config.yaml --segments 51 >& ../log_files/log.ace.inference.2020.2025.txt &

With RTX A5880 Ada 48 GB RAM, I found that I will run into out-of-memory issues by segment 13. I do not encounter such an issue with n_forward_steps=7543 when I run all inference steps at once.

[Z-Richard]

I have nailed down the memory to https://github.com/ai2cm/ace/blob/main/fme/ace/inference/inference.py#L263

stepper_config = config.load_stepper_config()

which, when running run_segmented_inference, will repeatedly load the model checkpoint into memory via torch.load(). An effective way to sidestep this problem without modifying the existing code base other than fme/ace/inference/inference.py would be to load stepper and stepper_config into memory first during main(), i.e.,

def main(
    yaml_config: str,
    segments: int | None = None,
    override_dotlist: Sequence[str] | None = None,
):
    config_data = prepare_config(yaml_config, override=override_dotlist)
    config = dacite.from_dict(
        data_class=InferenceConfig,
        data=config_data,
        config=dacite.Config(strict=True),
    )
    stepper = config.load_stepper()
    stepper.set_eval()
    stepper_config = config.load_stepper_config()
    prepare_directory(config.experiment_dir, config_data)
    if segments is None:
        with GlobalTimer():
            return run_inference_from_config(config, stepper, stepper_config)
    else:
        config.configure_logging(log_filename="inference_out.log")
        run_segmented_inference(config, stepper, stepper_config, segments)


def run_inference_from_config(config: InferenceConfig, stepper: Stepper, stepper_config: StepperConfig):
    timer = GlobalTimer.get_instance()
    timer.start_outer("inference")
    timer.start("initialization")

    if not os.path.isdir(config.experiment_dir):
        os.makedirs(config.experiment_dir, exist_ok=True)
    config.configure_logging(log_filename="inference_out.log")
    env_vars = logging_utils.retrieve_env_vars()
    beaker_url = logging_utils.log_beaker_url()
    config.configure_wandb(env_vars=env_vars, notes=beaker_url)

    if fme.using_gpu():
        torch.backends.cudnn.benchmark = True
    logging_utils.log_versions()
    logging.info(f"Current device is {fme.get_device()}")

    data_requirements = stepper_config.get_forcing_window_data_requirements(
        n_forward_steps=config.forward_steps_in_memory
    )

    logging.info("Loading initial condition data")
    initial_condition = get_initial_condition(
        config.initial_condition.get_dataset(), stepper_config.prognostic_names
    )

    logging.info("Initializing forcing data loader")
    print_vram("Before loading data")
    data = get_forcing_data(
        config=config.forcing_loader,
        total_forward_steps=config.n_forward_steps,
        window_requirements=data_requirements,
        initial_condition=initial_condition,
        surface_temperature_name=stepper.surface_temperature_name,
        ocean_fraction_name=stepper.ocean_fraction_name,
    )
    if not config.allow_incompatible_dataset:
        try:
            stepper.training_dataset_info.assert_compatible_with(data.dataset_info)
        except IncompatibleDatasetInfo as err:
            raise IncompatibleDatasetInfo(
                "Inference dataset is not compatible with dataset used for stepper "
                "training. Set allow_incompatible_dataset to True to ignore this "
                f"error. The incompatiblity found was: {str(err)}"
            ) from err
    variable_metadata = resolve_variable_metadata(
        dataset_metadata=data.variable_metadata,
        stepper_metadata=stepper.training_variable_metadata,
        stepper_all_names=stepper_config.all_names,
    )
    dataset_info = data.dataset_info.update_variable_metadata(variable_metadata)

    aggregator = config.aggregator.build(
        dataset_info=dataset_info,
        n_timesteps=config.n_forward_steps + stepper.n_ic_timesteps,
        output_dir=config.experiment_dir,
    )

    writer = config.get_data_writer(
        n_initial_conditions=data.n_initial_conditions,
        timestep=data.timestep,
        coords=data.coords,
        variable_metadata=variable_metadata,
    )

    timer.stop()
    logging.info("Starting inference")
    record_logs = get_record_to_wandb(label="inference")
    

    with torch.no_grad():
        run_inference(
            predict=stepper.predict_paired,
            data=data,
            writer=writer,
            aggregator=aggregator,
            record_logs=record_logs,
        )

    timer.start("final_writer_flush")
    logging.info("Starting final flush of data writer")
    writer.flush()
    logging.info("Writing reduced metrics to disk in netcdf format.")
    aggregator.flush_diagnostics()
    timer.stop()


    timer.stop_outer("inference")
    total_steps = config.n_forward_steps * data.n_initial_conditions
    inference_duration = timer.get_duration("inference")
    wandb_logging_duration = timer.get_duration("wandb_logging")
    total_steps_per_second = total_steps / (inference_duration - wandb_logging_duration)
    timer.log_durations()
    logging.info(
        "Total steps per second (ignoring wandb logging): "
        f"{total_steps_per_second:.2f} steps/second"
    )
    summary_logs = {
        "total_steps_per_second": total_steps_per_second,
        **timer.get_durations(),
        **aggregator.get_summary_logs(),
    }
    record_logs([summary_logs])

    del data, writer, stepper, stepper_config, aggregator, initial_condition, config
    del timer, env_vars, beaker_url, data_requirements
    gc.collect()
    torch.cuda.empty_cache()



def run_segmented_inference(config: InferenceConfig, stepper: Stepper, stepper_config: StepperConfig, segments: int):
    """Run inference in multiple segments.

    Args:
        config: inference configuration to be used for each individual segment. The
            provided initial condition configuration will only be used for the first
            segment.
        segments: total number of segments desired. Only missing segments will be run.

    Note:
        This is useful when running very long simulations or when saving a large
        amount of output data to disk. The simulation outputs will be split across
        multiple folders, each corresponding to one of the segments and labeled by
        the segment number.
    """
    logging.info(
        f"Starting segmented inference with {segments} segments. "
        f"Saving to {config.experiment_dir}."
    )
    config_copy = copy.deepcopy(config)
    original_wandb_name = os.environ.get("WANDB_NAME")
    for segment in range(segments):
        segment_label = f"segment_{segment:04d}"
        logging.info(f"Processing {segment_label}")
        segment_dir = os.path.join(config.experiment_dir, segment_label)
        restart_path = os.path.join(segment_dir, "restart.nc")
        if os.path.exists(restart_path):
            logging.info(f"Skipping segment {segment} because it has already been run.")
        else:
            logging.info(f"Running segment {segment}.")
            config_copy.experiment_dir = segment_dir
            if original_wandb_name is not None:
                os.environ["WANDB_NAME"] = f"{original_wandb_name}-{segment_label}"
            with GlobalTimer():
                run_inference_from_config(config_copy, stepper, stepper_config)
        config_copy.initial_condition = InitialConditionConfig(
            path=restart_path, engine="netcdf4"
        )

I know that --segments is meant to be an experimental feature and is subject to change, but would be happy to put together a PR to quickly fix this problem if needed.

[oliverwm1]

Thanks for the detailed report @Z-Richard. I don't think we'll want to change the API of run_inference_from_config but we could add garbage collection and/or torch cache emptying to the loop over segments in run_segmented_inference. Can you test if that resolves your issue?

[Z-Richard]

Hi @oliverwm1 I tried garbage collection but it did not solve the problem. The only thing that works for me is to construct stepper_config in main() and pass it to other functions.