ptxas spam warning

[ashep29]

When training GPT2, e.g., gpt2_causal_lm.fit(train_ds, epochs=num_epochs) ptxas warnings are spammed to output, e.g., messages similar to those below appear 100+ times in the output:

ptxas warning : Registers are spilled to local memory in function 'input_add_multiply_reduce_select_fusion_11', 64 bytes spill stores, 64 bytes spill loads

2025-03-23 20:38:42.249509: I external/local_xla/xla/stream_executor/cuda/cuda_asm_compiler.cc:397] ptxas warning : Registers are spilled to local memory in function 'gemm_fusion_dot_78_0', 56 bytes spill stores, 56 bytes spill loads

System info:
NAME="Rocky Linux"
VERSION="8.10 (Green Obsidian)"
ID="rocky"
ID_LIKE="rhel centos fedora"
VERSION_ID="8.10"
PLATFORM_ID="platform:el8"

[sonali-kumari1]

Hi @ashep29 -

Could you please share the sample code you're using to train GPT-2? Thanks!

[ashep29]

num_epochs = 3
chosen_preset = "gpt2_base_en"
batch_size=8

preprocessor = keras_nlp.models.GPT2CausalLMPreprocessor.from_preset(
chosen_preset,
sequence_length=128,
)

gpt2_causal_lm = keras_nlp.models.GPT2CausalLM.from_preset(
chosen_preset, preprocessor=preprocessor
)

---

Training data is a column from a pandas dataframe

only_perfect = df[df['rating'] == 5.0]
only_perfect.sample(5)

id 	review 	date 	rating

3878 croc_review_4305 I use this pair of Crocs as house slippers. Th... November 6, 2022 5.0
1772 croc_review_2338 I had a pair years ago that I literally wore o... July 13, 2023 5.0
820 croc_review_502 Comfy and cute! You can't go wrong with Crocs. March 27, 2021 5.0
2384 croc_review_2146 I find myself only wearing my Crocs now. I'm 7... January 9, 2022 5.0
7382 croc_review_7697 These are just like the original crocs that I ... July 8, 2022 5.0

tf_dataset = tf.data.Dataset.from_tensor_slices(only_perfect['review'].values)
train_ds = (tf_dataset.map(lambda review: review).batch(batch_size).prefetch(tf.data.AUTOTUNE))

learning_rate = keras.optimizers.schedules.PolynomialDecay(
5e-5,
decay_steps=train_ds.cardinality() * num_epochs,
end_learning_rate=0.0,
)

loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
optimiser = keras.optimizers.Adam(learning_rate)
metrics = [keras.metrics.SparseCategoricalAccuracy()]

gpt2_causal_lm.compile(
optimizer=optimiser,
loss=loss,
weighted_metrics=metrics,
)

gpt2_causal_lm.fit(train_ds, epochs=num_epochs)

---

Output:
Epoch 1/3

W0000 00:00:1742785503.537862 62622 assert_op.cc:38] Ignoring Assert operator compile_loss/sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/assert_equal_1/Assert/Assert
2025-03-24 14:05:10.179228: I external/local_xla/xla/stream_executor/cuda/cuda_asm_compiler.cc:397] ptxas warning : Registers are spilled to local memory in function 'gemm_fusion_dot_280', 76 bytes spill stores, 76 bytes spill loads

2025-03-24 14:05:10.214570: I external/local_xla/xla/stream_executor/cuda/cuda_asm_compiler.cc:397] ptxas warning : Registers are spilled to local memory in function 'gemm_fusion_dot_280', 4 bytes spill stores, 4 bytes spill loads

2025-03-24 14:05:10.484047: I external/local_xla/xla/stream_executor/cuda/cuda_asm_compiler.cc:397] ptxas warning : Registers are spilled to local memory in function 'gemm_fusion_dot_280_0', 376 bytes spill stores, 328 bytes spill loads

2025-03-24 14:05:10.690244: I external/local_xla/xla/stream_executor/cuda/cuda_asm_compiler.cc:397] ptxas warning : Registers are spilled to local memory in function 'gemm_fusion_dot_280', 40 bytes spill stores, 40 bytes spill loads

2025-03-24 14:05:10.721397: I external/local_xla/xla/stream_executor/cuda/cuda_asm_compiler.cc:397] ptxas warning : Registers are spilled to local memory in function 'gemm_fusion_dot_280', 332 bytes spill stores, 336 bytes spill loads

2025-03-24 14:05:10.756273: I external/local_xla/xla/stream_executor/cuda/cuda_asm_compiler.cc:397] ptxas warning : Registers are spilled to local memory in function 'gemm_fusion_dot_76_0', 1144 bytes spill stores, 1348 bytes spill loads

2025-03-24 14:05:10.987391: I external/local_xla/xla/stream_executor/cuda/cuda_asm_compiler.cc:397] ptxas warning : Registers are spilled to local memory in function 'gemm_fusion_dot_33618', 56 bytes spill stores, 56 bytes spill loads

....this is repeated hundreds of times

[ashep29]

The ptxas warning also appears (not spammed as often, but still a pain) during inference before any training is implemented, e.g.,

preprocessor = keras_nlp.models.GPT2CausalLMPreprocessor.from_preset(
chosen_preset,
sequence_length=128,
)

gpt2_causal_lm = keras_nlp.models.GPT2CausalLM.from_preset(
chosen_preset, preprocessor=preprocessor
)

sentence_starters = ["I bought a pair of crocs", "Crocs are","Received a pair of crocs"]
max_length = 75

output = gpt2_causal_lm.generate(sentence_starters,max_length=max_length)

for sentence in output:
print(f"--{sentence}\n")

2025-03-31 13:17:53.430428: E tensorflow/core/util/util.cc:131] oneDNN supports DT_INT64 only on platforms with AVX-512. Falling back to the default Eigen-based implementation if present.
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1743387476.255537 3897003 service.cc:148] XLA service 0x5626c18924d0 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
I0000 00:00:1743387476.255570 3897003 service.cc:156] StreamExecutor device (0): NVIDIA A100-PCIE-40GB, Compute Capability 8.0
2025-03-31 13:17:56.352019: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:268] disabling MLIR crash reproducer, set env var MLIR_CRASH_REPRODUCER_DIRECTORY to enable.
I0000 00:00:1743387477.001111 3897003 cuda_dnn.cc:529] Loaded cuDNN version 90800
2025-03-31 13:17:59.044107: I external/local_xla/xla/stream_executor/cuda/cuda_asm_compiler.cc:397] ptxas warning : Registers are spilled to local memory in function 'gemm_fusion_dot_399', 8 bytes spill stores, 8 bytes spill loads

2025-03-31 13:17:59.215341: I external/local_xla/xla/stream_executor/cuda/cuda_asm_compiler.cc:397] ptxas warning : Registers are spilled to local memory in function 'gemm_fusion_dot_399', 80 bytes spill stores, 80 bytes spill loads

2025-03-31 13:17:59.242393: I external/local_xla/xla/stream_executor/cuda/cuda_asm_compiler.cc:397] ptxas warning : Registers are spilled to local memory in function 'gemm_fusion_dot_420', 360 bytes spill stores, 328 bytes spill loads

2025-03-31 13:17:59.267403: I external/local_xla/xla/stream_executor/cuda/cuda_asm_compiler.cc:397] ptxas warning : Registers are spilled to local memory in function 'gemm_fusion_dot_399', 360 bytes spill stores, 328 bytes spill loads

2025-03-31 13:17:59.620412: I external/local_xla/xla/stream_executor/cuda/cuda_asm_compiler.cc:397] ptxas warning : Registers are spilled to local memory in function 'gemm_fusion_dot_420', 80 bytes spill stores, 80 bytes spill loads

2025-03-31 13:17:59.934400: I external/local_xla/xla/stream_executor/cuda/cuda_asm_compiler.cc:397] ptxas warning : Registers are spilled to local memory in function 'gemm_fusion_dot_399', 20 bytes spill stores, 20 bytes spill loads

2025-03-31 13:18:00.547474: I external/local_xla/xla/stream_executor/cuda/cuda_asm_compiler.cc:397] ptxas warning : Registers are spilled to local memory in function 'gemm_fusion_dot_420', 8 bytes spill stores, 8 bytes spill loads

2025-03-31 13:18:00.727426: I external/local_xla/xla/stream_executor/cuda/cuda_asm_compiler.cc:397] ptxas warning : Registers are spilled to local memory in function 'gemm_fusion_dot_420', 20 bytes spill stores, 20 bytes spill loads

2025-03-31 13:18:02.064344: I external/local_xla/xla/stream_executor/cuda/cuda_asm_compiler.cc:397] ptxas warning : Registers are spilled to local memory in function 'gemm_fusion_dot_399_0', 28328 bytes spill stores, 31092 bytes spill loads

2025-03-31 13:18:02.264268: I external/local_xla/xla/stream_executor/cuda/cuda_asm_compiler.cc:397] ptxas warning : Registers are spilled to local memory in function 'gemm_fusion_dot_420_0', 28328 bytes spill stores, 31092 bytes spill loads

--I bought a pair of crocs from the local supermarket last year, and the one with the red lace was so good that I thought I'd buy them.

I'm not really into lace, but this is the perfect size for a crochet, so I'm happy to have them in stock. They're not too small at all - I'd say they

--Crocs are the most common form of skin in the world.

The Croc is an extremely rare skin condition that affects almost every type of skin in the world. Crocs are often referred to as "the skin condition." They are caused by inflammation of the skin cells.

Crocs can cause skin conditions such as acne, psoriasis, ec

--Received a pair of crocs. Very nice.

I'm a very happy camper. I have a couple of dogs. They were very happy and I was very excited.

My wife is a little older and has a lot of trouble with her hands. I'm happy with them.

Thank you, and happy camper.

I0000 00:00:1743387486.470904 3897003 device_compiler.h:188] Compiled cluster using XLA! This line is logged at most once for the lifetime of the process.

[sonali-kumari1]

Hi @ashep29 -

keras_nlp has been renamed to keras_hub, so you can install keras_hub using !pip install -Uq keras-huband I am not getting any ptxas spam warnings on Colab notebook with the latest version of keras(3.9.1) and keras_hub(0.19.3). Attaching gist for your reference. Thanks!

[github-actions]

This issue is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you.