Run on Apple Silicon GPUs

[dragosmc]

Hi.

I'm running ludwig locally on an Apple Silicon and I was wondering what should I change to run with PyTorch on MPS.
The provided instructions on ludwig.ai are for CUDA only and couldn't find anything but this patch master...rudolfolah:ludwig:patch-2 for enabling somehow "mps".

I am not that familiar with the architecture of ludwig and pytorch but I'm happy to debug/lend a hand in getting MPS to work.

Dragos.

[tgaddair]

Hey @dragosmc, were you able to try the patch linked above? That was going to be my suggestion to try first. Unfortunately, I'm running on an Intel Mac and can't test directly, but happy to help debug if you run into issues with the patch. If it works, let's add it to Ludwig officially.

[dragosmc]

Hey, I've tried this fix out and I have some feedback for a couple of scenarios. Before that the specs I'm running this on are:

• M2 Pro with 12‑core CPU, 19‑core GPU and 16‑core Neural Engine
• Python 3.8
• Latest installed PyTorch by following the Conda instructions here https://pytorch.org/get-started/locally/
• changed the get_torch_device function as described in the mentioned PR above

pytorch-1.13.1             |          py3.8_0        44.1 MB  pytorch
torchaudio-0.13.1          |         py38_cpu         5.6 MB  pytorch
torchvision-0.14.1         |         py38_cpu         6.2 MB  pytorch

I have the following test setup:

• one set of tabular features (3 numerical columns for a number of rows)
• one sequential feature

Upon running the above setup I get the error below.

NotImplementedError: The operator 'aten::sort.values_stable' is not currently implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on https://github.com/pytorch/pytorch/issues/77764. As a temporary fix, you can set the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.

After following instructions I was able to successfully run the tabular data only and it ran on GPU as I could see the GPU usage histogram from Activity Monitor. The same happened with running with both tabular and sequential features with the default parameters for the sequential one.

I then changed the sequential feature to use the following setup and the train process stuck:

{
	'type': 'sequence',
	'name': col,
	'encoder': {
		'type': 'rnn',
		'embedding_size': 64,
		'cell_type': 'lstm',
		'num_layers': 3
	}
}

To recap (checkmark denotes a successful experiment):

• Tried running on GPU with tabular features only
• Tried running on GPU with tabular features and sequential features with default config
• Tried running on GPU with tabular features and sequential features with the rnn type
• Tried running on GPU with tabular features and sequential features with transform type

I'm happy to help out with more debugging.

[justinxzhao]

This code in embedding_modules.py looks like it might need to be updated.

Numerical features don't call the embeddings module, which explains why you didn't run into any issues in tabular tests. What was the last test you ran with transform?

[dragosmc]

Hey @justinxzhao, thanks for that pointer to the embedding_modules. Regarding the test was a simple sequence for which I changed the configuration of the input from type: rnn to type: transform.

[tgaddair]

Thanks for looking into this, @dragosmc! We just landed #3072 and #3074 to allow users to use experimental MPS support when LUDWIG_ENABLE_MPS=1 is set in the environment. It sounds like your experience is similar to @arnavgarg1 's when he was testing this out.

Here were the results we saw:

WITH MPS

• Tabular with defaults - 45 epochs - 19s
• Tabular with tabnet + defaults - 9 epochs - 20s
• Text (1000 rows) with defaults - 8 epochs - 16s
• Text (1000 rows) with BERT encoder - 7 epochs - 108s

WITHOUT MPS

• Tabular with defaults - 45 epochs - 6s
• Tabular with tabnet + defaults - 9 epochs - 5s
• Text (1000 rows) with defaults - 6 epochs - 38s
• Text (1000 rows) with BERT encoder - cancelled, but ~ 4 minutes per epoch

In all cases we also had to set PYTORCH_ENABLE_MPS_FALLBACK=1 as some ops weren't supported. But seems like most things are working with the fallback var set, except RNN and possibly a few other things we haven't tested yet.

I'm curious @dragosmc, did you notice any significant performance improvement or degradation in your tests?

[arnavgarg1]

Just for visibility, here's some thoughts I had on performance difference for tabular datasets:

• Since MPS technically falls back to CPU on occasion, I could see CPU utilization was fairly high for tabular, but very low for text (~40% usage). In contrast, when using the CPU device, I noticed that CPU utilization was even higher for tabular (~85%), and significantly higher for text purely on CPU (around 95% usage).
• I wonder how much of a difference regression vs classification tasks make. I noticed more warnings when doing regression (when computing roc_auc for e.g. because it had to fall back to using the CPU) compared to classification, so maybe it spent more time on the MPS device in the latter case and that's where the context switch + data transfer overhead comes from for tabular datasets (mine was a regression problem that I was benchmarking on)