diff --git a/utils/model_ops/utils/torchop_yaml.py b/utils/model_ops/utils/torchop_yaml.py index 86d8e795..376d8306 100644 --- a/utils/model_ops/utils/torchop_yaml.py +++ b/utils/model_ops/utils/torchop_yaml.py @@ -181,12 +181,19 @@ def sanitize_arg( return str(arg) +# Op-name prefixes that never denote a tensor operation, so they get no test case. +# torch.cuda.* are host-side device queries (e.g. get_device_capability(), called +# from transformers' _can_use_grouped_mm) that return Python values, not tensors. +_SKIPPED_OP_PREFIXES = ("torch.cuda.",) + # Ops whose tensor inputs should use Xavier init when dtype/rank also qualify. _XAVIER_OPS = { "torch.conv2d", "torch.bmm", "torch.matmul", + "torch._grouped_mm", "torch.nn.functional.linear", + "torch.nn.functional.grouped_mm", } _XAVIER_DTYPES = {"torch.float16", "torch.float32", "torch.bfloat16"} @@ -221,32 +228,85 @@ def _maybe(t): _maybe(t) +# Grouped-GEMM group-boundary kwargs: (op, kwarg name) -> positional input whose +# leading dim gives the total row count that the offsets must end at. +_CUMSUM_OFFSETS_KWARGS = { + ("torch._grouped_mm", "offs"): 0, + ("torch.nn.functional.grouped_mm", "offs"): 0, + ("torch.ops.transformers.grouped_mm_fallback", "offs"): 0, +} + + +def _maybe_promote_init_to_cumsum_offsets(op_name, kwarg_name, details, yaml_inputs): + """Flip a grouped-GEMM ``offs`` tensor from the generic integer ``init: randint`` + to ``init: cumsum_offsets``. + + ``offs`` is not an arbitrary integer tensor: the kernel reads it as cumulative + group boundaries, so it must be monotonic non-decreasing and end at the number + of rows in ``mat_a`` (see transformers' ``grouped_mm_experts_forward``, which + builds it as ``cumsum(histc(sort(expert_ids)))``). Random values corrupt the + group boundaries, so record the constraint instead: ``init_args.total`` carries + ``mat_a.size(0)``, taken from the already-formatted positional inputs. + Mutates ``details`` in place. + """ + arg_index = _CUMSUM_OFFSETS_KWARGS.get((op_name, kwarg_name)) + if arg_index is None or not isinstance(details, dict): + return + if "int" not in str(details.get("dtype")): + return + shape = details.get("shape") + if shape is None or len(shape) != 1: + return + + total = None + if yaml_inputs and len(yaml_inputs) > arg_index: + inp = yaml_inputs[arg_index] + if isinstance(inp, dict) and isinstance(inp.get("tensor"), dict): + mat_a_shape = inp["tensor"].get("shape") + if mat_a_shape is not None and len(mat_a_shape) >= 1: + total = int(mat_a_shape[0]) + if total is None: + return + + details["init"] = "cumsum_offsets" + details.pop("init_args", None) + details["init_args"] = {"total": total} + + +def build_tensor_details(shape, stride, storage_offset, dtype, device, randintlim=1000): + """Build the YAML tensor block (shape/stride/dtype/device/init) for a tensor + described purely by its metadata.""" + tensor_details = {} + tensor_details["shape"] = FlowList(shape) + tensor_details["stride"] = FlowList(stride) + tensor_details["storage_offset"] = storage_offset + tensor_details["dtype"] = str(dtype) + tensor_details["device"] = str(device) + if "float" in str(dtype): + tensor_details["init"] = "rand" + elif "int" in str(dtype): + tensor_details["init"] = "randint" + tensor_details["init_args"] = {"high": randintlim} + elif "bool" in str(dtype): + pass + else: + print(f"Unhandled dtype {dtype}") + return {} + + return tensor_details + + def format_tensor_details( arg, shape, stride, storage_offset, dtype, device, randintlim=1000 ): logging.getLogger("TorchOpCollector").debug( f"In format_tensor_details: {arg}, {shape}, {stride}, {storage_offset}, {dtype}, {device}, {randintlim}" ) - tensor_details = {} if isinstance(arg, torch.fx.node.Node): - tensor_details["shape"] = FlowList(shape) - tensor_details["stride"] = FlowList(stride) - tensor_details["storage_offset"] = storage_offset - tensor_details["dtype"] = str(dtype) - tensor_details["device"] = str(device) - if "float" in str(dtype): - tensor_details["init"] = "rand" - elif "int" in str(dtype): - tensor_details["init"] = "randint" - tensor_details["init_args"] = {"high": randintlim} - elif "bool" in str(dtype): - pass - else: - print(f"Unhandled dtype {dtype}") - return {} - - return tensor_details + return build_tensor_details( + shape, stride, storage_offset, dtype, device, randintlim + ) def add_test_case_yaml( @@ -891,15 +951,24 @@ def _pack_reshape_args( TorchOpCollector.log_function[TorchOpCollector.log_mthd]( f"Converting scalar args {san_args} to tuple for {op_name}" ) - start_index_norm = 1 + # The shape/dims may arrive either as varargs -- reshape(t, 1, 24, 2816) -- + # or as a single sequence -- reshape(t, (1, 24, 2816)). Flatten both to a + # plain list of dims before indexing. + dims = [] + for arg in san_args[1:]: + if isinstance(arg, (tuple, list, torch.Size)): + dims.extend(arg) + else: + dims.append(arg) + start_index_norm = 0 if dim_redn: TorchOpCollector.log_function[TorchOpCollector.log_mthd]( f"Dim reduction: {orig_shape}, {shrunk_shape}, {start_index_norm}" ) dims_removed = len(orig_shape) - len(shrunk_shape) - start_index_norm = max(dims_removed, 0) + 1 - val_tuple = tuple(int(ii) for ii in san_args[1:]) - val_tuple_norm = tuple(int(ii) for ii in san_args[start_index_norm:]) + start_index_norm = max(dims_removed, 0) + val_tuple = tuple(int(ii) for ii in dims) + val_tuple_norm = tuple(int(ii) for ii in dims[start_index_norm:]) if op_name == "torch.reshape": if round_up: TorchOpCollector.log_function[TorchOpCollector.log_mthd]( @@ -1021,7 +1090,10 @@ def _emit_test_case( TorchOpCollector.log_function[TorchOpCollector.log_mthd]( "Creating round_up version of the yaml" ) - kwmap_norm = copy.deepcopy(kwmap) + # Shallow copy: only scalar entries ("dtype"/"fill_value") are + # rewritten below, and a deep copy would recurse into non-copyable + # values such as FakeTensors reachable from fx Nodes. + kwmap_norm = copy.copy(kwmap) if ( len(kwmap_norm) > 0 and "dtype" in kwmap_norm @@ -1119,6 +1191,11 @@ def collect_torchops(gm, ops_set, print_output): f"Skipping op {op_name} as its shape contains unbound symbolic elements" ) continue + if op_name.startswith(_SKIPPED_OP_PREFIXES): + TorchOpCollector.log_function[TorchOpCollector.log_mthd]( + f"Skipping op {op_name} as it is not a tensor operation" + ) + continue TorchOpCollector.log_function[TorchOpCollector.log_mthd]( f"MAIN OP DETAILS: {op_name}, {out_shape}, {out_stride}, {out_storage_offset}, {out_dtype}, {out_device}" ) @@ -1219,6 +1296,28 @@ def collect_torchops(gm, ops_set, print_output): kwmap[k] = v if isinstance(v, torch.dtype) or k == "device" or v is None: kwmap[k] = str(v) + elif isinstance(v, torch.fx.Node): + # e.g. torch._grouped_mm(..., offs=). + # Record the tensor's shape/dtype instead of the + # producing variable's name: an fx Node carries a + # FakeTensor in node.meta, which is neither + # YAML-serializable nor deep-copyable (data_ptr() raises). + details = TorchOpCollector._tensor_kwarg_details(v) + _maybe_promote_init_to_cumsum_offsets( + op_name, k, details, yaml_inputs + ) + # Nest under "tensor" so a tensor-valued kwarg is + # described exactly like a positional tensor input. + kwmap[k] = ( + {"tensor": details} if details is not None else str(v) + ) + elif isinstance(v, torch.Tensor): + # A constant tensor baked into the graph: describe it + # the same way, so kwmap never holds a live tensor. + details = TorchOpCollector._tensor_details_from_meta(v) + kwmap[k] = ( + {"tensor": details} if details is not None else str(v) + ) output_list.append( f" output shape: {out_shape}, {out_stride}, {out_storage_offset}, {out_dtype}, {out_device} @ {node.name}" @@ -1275,6 +1374,38 @@ def _extract_stacktrace_comments(node): comments = comments[comm_idx:-2] return comments, stacktrace + @staticmethod + def _tensor_kwarg_details(node): + """Describe a tensor-valued kwarg (e.g. torch._grouped_mm's `offs`) by its + shape/dtype rather than by the producing variable's name. + + Only metadata is available here: the collector runs at compile time, where + the kwarg's value is a FakeTensor with no backing data (reading it raises + DataDependentOutputException for data-dependent producers such as + cumsum(histc(...))). Emit the same tensor block used for positional + tensor inputs so the test harness can synthesize an input of the right + shape and dtype. + """ + meta_val = node.meta.get( + "val", + node.meta.get("tensor_meta", node.meta.get("example_value", None)), + ) + if meta_val is None: + return None + return TorchOpCollector._tensor_details_from_meta(meta_val) + + @staticmethod + def _tensor_details_from_meta(meta_val): + """Build a YAML tensor block from a (fake) tensor's metadata.""" + shape, stride, storage_offset, dtype, device = ( + TorchOpCollector._extract_meta_info(meta_val) + ) + if shape is None or dtype is None or shape == "symbolic": + return None + return ( + build_tensor_details(shape, stride, storage_offset, dtype, device) or None + ) + @staticmethod def _extract_meta_info(meta_val): if ( @@ -1363,6 +1494,11 @@ def _extract_meta_info(meta_val): dtypes_tuple = () devices_tuple = () for vv in meta_val: + if not isinstance(vv, torch.Tensor): + TorchOpCollector.log_function[TorchOpCollector.log_mthd]( + f"{vv}: {type(vv)}" + ) + continue TorchOpCollector.log_function[TorchOpCollector.log_mthd]( f"{vv}: {type(vv)} type: {vv.dtype} shape: {vv.shape}" )