_load_for_executorch pybinding cannot load model for training

[jackzhxng]

🐛 Describe the bug

_load_for_executorch pybinding cannot load joint graph for phi-3-mini-lora because load_into is not implemented in MmapLoader, and load_into is used by Program's load_mutable_subsegment_into to load weights for training in TensorParser.

Versions

Collecting environment information...
PyTorch version: 2.5.0.dev20240821
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.6.1 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.3.9.4)
CMake version: version 3.30.0-rc3
Libc version: N/A

Python version: 3.10.0 (default, Mar 3 2022, 03:54:28) [Clang 12.0.0 ] (64-bit runtime)
Python platform: macOS-14.6.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M1 Pro

Versions of relevant libraries:
[pip3] executorch==0.4.0a0+3ad2959
[pip3] numpy==1.21.3
[pip3] torch==2.5.0.dev20240821
[pip3] torchaudio==2.4.0.dev20240821
[pip3] torchsr==1.0.4
[pip3] torchvision==0.20.0.dev20240821
[conda] executorch 0.4.0a0+3ad2959 pypi_0 pypi
[conda] numpy 1.21.3 pypi_0 pypi
[conda] torch 2.5.0.dev20240821 pypi_0 pypi
[conda] torchaudio 2.4.0.dev20240821 pypi_0 pypi
[conda] torchsr 1.0.4 pypi_0 pypi
[conda] torchvision 0.20.0.dev20240821 pypi_0 pypi

cc @JacobSzwejbka @dbort

[JacobSzwejbka]

This is because the mmap_data_loader does not currently implement load_into.

[digantdesai]

close this?

[jackzhxng]

load_into is still not implemented for mmap data loader. Not a bug so moving to feature request

[dbort]

Added good first issue because this should be pretty straightforward.

• https://github.com/pytorch/executorch/blob/main/extension/data_loader/mmap_data_loader.h is the class to fix.

• executorch/runtime/core/data_loader.h

  Lines 111 to 115 in d99970b

  	ET_NODISCARD virtual Error load_into(
  	size_t offset,
  	size_t size,
  	const SegmentInfo& segment_info,
  	void* buffer) const {

  is the missing virtual method that MmapDataLoader should override.
  • Although we could mmap and then memcpy the data, that's pretty wasteful. It'd be better to do exactly the same thing that FileDataLoader::load_into does: read from the file descriptor directly into the provided buffer. We could refactor that FileDataLoader load_into logic into a common helper shared by both DataLoaders.