From a2ea5966df9486814541414b90bd755a3d499366 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Fri, 31 Oct 2025 10:08:59 +0000 Subject: [PATCH] Allow XLA memory fraction greater than one when using unified memory --- src/xla/XLA.jl | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/xla/XLA.jl b/src/xla/XLA.jl index f14139b890..b6b9148a86 100644 --- a/src/xla/XLA.jl +++ b/src/xla/XLA.jl @@ -134,8 +134,13 @@ function __init__() ) @debug "XLA_REACTANT_GPU_MEM_FRACTION: " XLA_REACTANT_GPU_MEM_FRACTION[] maxlog = 1 - if XLA_REACTANT_GPU_MEM_FRACTION[] > 1 || XLA_REACTANT_GPU_MEM_FRACTION[] < 0 - error("XLA_REACTANT_GPU_MEM_FRACTION must be between 0 and 1") + if XLA_REACTANT_GPU_MEM_FRACTION[] < 0 + error("XLA_REACTANT_GPU_MEM_FRACTION must be not be negative") + elseif XLA_REACTANT_GPU_MEM_FRACTION[] > 1 + if get(ENV, "TF_FORCE_UNIFIED_MEMORY", "0") != "1" + error("XLA_REACTANT_GPU_MEM_FRACTION must be not greater than 1 without \ + TF_FORCE_UNIFIED_MEMORY=1") + end end end